diff --git a/socks5.go b/socks5.go index 22aa07b..01698ea 100644 --- a/socks5.go +++ b/socks5.go @@ -140,13 +140,13 @@ func socks5Handle(conn net.Conn) { req, err := gosocks5.ReadRequest(conn) if err != nil { - //log.Println(err) + log.Println(err) return } switch req.Cmd { case gosocks5.CmdConnect: - //log.Println("connect", req.Addr.String()) + log.Println("connect", req.Addr.String()) tconn, err := Connect(req.Addr.String(), Proxy) if err != nil { gosocks5.NewReply(gosocks5.HostUnreachable, nil).Write(conn) diff --git a/ws.go b/ws.go index 017fd4b..0f7522a 100644 --- a/ws.go +++ b/ws.go @@ -20,16 +20,20 @@ func NewWSConn(conn *websocket.Conn) *WSConn { } func (conn *WSConn) Read(b []byte) (n int, err error) { - _, b, err = conn.ReadMessage() - n = len(b) - - return + _, r, err := conn.NextReader() + if err != nil { + return + } + return r.Read(b) } func (conn *WSConn) Write(b []byte) (n int, err error) { - n = len(b) - err = conn.WriteMessage(websocket.BinaryMessage, b) - return + w, err := conn.NextWriter(websocket.BinaryMessage) + if err != nil { + return + } + defer w.Close() + return w.Write(b) } func (conn *WSConn) SetDeadline(t time.Time) error { @@ -58,9 +62,12 @@ func (s *WSServer) handle(w http.ResponseWriter, r *http.Request) { } defer conn.Close() - c := NewWSConn(conn) - - socks5Handle(gosocks5.ServerConn(c, serverConfig)) + c := gosocks5.ServerConn(NewWSConn(conn), serverConfig) + if err := c.Handleshake(); err != nil { + log.Println(err) + return + } + socks5Handle(c) } func (s *WSServer) ListenAndServe() error {