package net import ( "b612.me/starlog" "bytes" "errors" "fmt" "io" "net" "strings" "sync" "time" ) var MSG_CMD_HELLO = []byte{11, 27, 19, 96, 182, 18, 25, 150, 17, 39} var MSG_NEW_CONN = []byte{0, 0, 0, 0, 255, 255, 255, 255, 11, 27} var MSG_NEW_CONN_REQ = []byte{0, 0, 0, 0, 255, 255, 255, 255, 19, 96} var MSG_CLOSE = []byte{255, 255, 0, 0, 255, 0, 0, 255, 255, 27} var MSG_HEARTBEAT = []byte{6, 66, 66, 6, 6, 66, 6, 66, 11, 27} type SimpleNatServer struct { mu sync.RWMutex cmdTCPConn net.Conn cmdUDPConn *net.UDPAddr listenTcp net.Listener listenUDP *net.UDPConn Addr string Port int lastTCPHeart int64 lastUDPHeart int64 Passwd string DialTimeout int64 UDPTimeout int64 running int32 tcpConnPool chan net.Conn tcpAlived bool } func (s *SimpleNatServer) getConnfromTCPPool() (net.Conn, error) { select { case conn := <-s.tcpConnPool: return conn, nil case <-time.After(time.Second * 10): return nil, errors.New("no connection got") } } func (s *SimpleNatServer) tcpCmdConn() net.Conn { s.mu.RLock() defer s.mu.RUnlock() return s.cmdTCPConn } func (s *SimpleNatServer) tcpCmdConnAlived() bool { s.mu.RLock() defer s.mu.RUnlock() return s.tcpAlived } func (s *SimpleNatServer) listenTCP() error { var err error s.tcpConnPool = make(chan net.Conn, 10) s.listenTcp, err = net.Listen("tcp", fmt.Sprintf("%s:d", s.Addr, s.Port)) if err != nil { starlog.Errorln("failed to listen tcp", err) return err } for { conn, err := s.listenTcp.Accept() if err != nil { continue } if s.tcpCmdConnAlived() { go s.tcpClientServe(conn.(*net.TCPConn)) continue } go s.waitingForTCPCmd(conn.(*net.TCPConn)) } return nil } func (s *SimpleNatServer) tcpClientServe(conn *net.TCPConn) { if !s.tcpCmdConnAlived() { conn.Close() return } if strings.Split(conn.RemoteAddr().String(), ":")[0] == strings.Split(s.tcpCmdConn().RemoteAddr().String(), ":")[0] { conn.SetReadDeadline(time.Now().Add(5 * time.Second)) cmdBuf := make([]byte, 10) if _, err := io.ReadFull(conn, cmdBuf); err == nil { conn.SetReadDeadline(time.Time{}) if bytes.Equal(cmdBuf, MSG_NEW_CONN) { starlog.Noticef("Nat Server Recv New Client Conn From %v\n", conn.RemoteAddr().String()) s.tcpConnPool <- conn return } } conn.SetReadDeadline(time.Time{}) } starlog.Noticef("Nat Server Recv New Side Conn From %v\n", conn.RemoteAddr().String()) _, err := s.tcpCmdConn().Write(MSG_NEW_CONN_REQ) if err != nil { s.mu.Lock() s.cmdTCPConn.Close() s.tcpAlived = false s.mu.Unlock() starlog.Errorf("Failed to Write CMD To Client:%v\n", err) return } reverse, err := s.getConnfromTCPPool() if err != nil { starlog.Errorf("Nat Server Conn to %v Closed %v\n", conn.RemoteAddr(), err) conn.Close() return } starlog.Infof("Nat Server Conn %v<==>%v Connected\n", conn.RemoteAddr(), reverse.RemoteAddr()) Copy(reverse, conn) starlog.Warningf("Nat Server Conn %v<==>%v Closed\n", conn.RemoteAddr(), reverse.RemoteAddr()) } func (s *SimpleNatServer) waitingForTCPCmd(conn *net.TCPConn) { conn.SetReadDeadline(time.Now().Add(time.Duration(s.DialTimeout) * time.Second)) cmdBuf := make([]byte, 10) if _, err := io.ReadFull(conn, cmdBuf); err != nil { conn.Close() return } if bytes.Equal(cmdBuf, MSG_CMD_HELLO) { s.mu.Lock() s.cmdTCPConn = conn s.tcpAlived = true conn.SetKeepAlive(true) conn.SetKeepAlivePeriod(time.Second * 20) s.mu.Unlock() } }