Socks5 server in single Go file

rockem-sockem.go

package main

import (
        "log"
        "net"
        "strconv"
)

type connection interface {
        Bind(conn net.Conn)

        Close(reason uint8)

        negotiate()
}

type SocksConn struct {
        conn net.Conn
}

func main() {
        l, err := net.Listen("tcp", ":8888")
        if err != nil {
                log.Fatal("bind port error.")
                return
        }
        log.Println("server start, listening 8888")
        for {
                c, err := l.Accept()
                if err != nil {
                        log.Fatal("cannot accept client")
                        continue
                }
                go handle(c)
        }
}

func handle(conn net.Conn) {
        socks := new(SocksConn)
        socks.Bind(conn)
}

func (s *SocksConn) Bind(conn net.Conn) {
        s.conn = conn
        buf := make([]byte, 3, 256)
        _, err := conn.Read(buf)
        log.Printf("new connection incoming, len:%v, cap:%v", len(buf), len(buf))
        version := buf[0]
        if version != 0x05 {
                s.Close(0x80)
        }
        methodLen := buf[1]
        rsp := make([]byte, 2)
        rsp[0] = 0x05
        method, err := selectMethod(buf[2 : 2+int(methodLen)])
        if err != nil {
                s.Close(0xff)
                return
        }
        rsp[1] = method
        _, err = conn.Write(rsp)
        if err != nil {
                s.Close(0x80)
                return
        }
        s.negotiate()
}

//TODO implements cmd: BIND, UDP ASSOCIATE
func (s *SocksConn) negotiate() {
        log.Println("negotiating with client")
        buf := make([]byte, 5)
        if _, err := s.conn.Read(buf); err != nil {
                s.Close(0x80)
                return
        }
        if version := buf[0]; version != 0x05 {
                s.Close(0x80)
                return
        }
        //cmd := buf[1]
        addrType := buf[3]
        var addrSize byte
        switch addrType {
        case 0x01:
                addrSize = 4
        case 0x03:
                addrSize = buf[4] + 1
        case 0x04:
                addrSize = 16
        }
        follow := make([]byte, addrSize+1)
        if _, err := s.conn.Read(follow); err != nil {
                s.Close(0x80)
                return
        }
        buf = append(buf, follow...)
        //TODO IPv6
        addr, err := net.ResolveTCPAddr("tcp", getAddr(buf[3:]))
        if err != nil {
                s.Close(0x03)
                return
        }
        log.Println("trying", addr)
        req, e := net.DialTCP("tcp", nil, addr)
        if e != nil {
                log.Println(e)
                s.Close(0x03)
                return
        }
        buf[1] = 0x00
        if _, err = s.conn.Write(buf); err != nil {
                log.Println(err)
                s.Close(0x80)
                return
        }
        go listenDownstream(*s, req)
        go listenUpstream(*s, req)
}

func (s *SocksConn) Close(reason uint8) {
        if reason != 0x80 {
                buf := make([]byte, 2)
                buf[0] = 0x05
                buf[1] = 0xff
                s.conn.Write(buf)
        }
        s.conn.Close()
        log.Printf("connection closed, code:%d", reason)
}

func listenDownstream(down SocksConn, up net.Conn) {
        defer down.Close(0x80)
        defer up.Close()
        buf := make([]byte, 65536)
        for {
                size, err := down.conn.Read(buf)
                if err != nil {
                        return
                }
                _, err = up.Write(buf[0:size])
                if err != nil {
                        return
                }
        }
}

func listenUpstream(down SocksConn, up net.Conn) {
        defer down.Close(0x03)
        defer up.Close()
        buf := make([]byte, 65536)
        for {
                size, err := up.Read(buf)
                if err != nil {
                        return
                }
                _, err = down.conn.Write(buf[0:size])
                if err != nil {
                        return
                }
        }
}

//TODO
func selectMethod(methods []uint8) (uint8, error) {
        return 0x00, nil
}

func getAddr(b []byte) string {
        var addr string
        switch b[0] {
        case 0x01:
                addr = strconv.Itoa(int(b[1])) + "." + strconv.Itoa(int(b[2])) + "." + strconv.Itoa(int(b[3])) + "." + strconv.Itoa(int(b[4]))
        case 0x03:
                addr = string(b[2 : len(b)-2])
                //TODO IPv6
        }
        port := int(b[len(b)-2])<<8 + int(b[len(b)-1])
        addr += ":" + strconv.Itoa(port)
        return addr
}