Merge pull request txthinking#1 from onoketa/master

Use more generic io interface
jinrenjie · Mar 27, 2020 · 9f204fc · 9f204fc
2 parents 7d82196 + 6506f90
commit 9f204fc
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/client_side.go b/client_side.go
@@ -4,7 +4,6 @@ import (
 	"errors"
 	"io"
 	"log"
-	"net"
 )
 
 var (
@@ -22,7 +21,7 @@ func NewNegotiationRequest(methods []byte) *NegotiationRequest {
 }
 
 // WriteTo write negotiation request packet into server
-func (r *NegotiationRequest) WriteTo(w *net.TCPConn) error {
+func (r *NegotiationRequest) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver}); err != nil {
 		return err
 	}
@@ -39,7 +38,7 @@ func (r *NegotiationRequest) WriteTo(w *net.TCPConn) error {
 }
 
 // NewNegotiationReplyFrom read negotiation reply packet from server
-func NewNegotiationReplyFrom(r *net.TCPConn) (*NegotiationReply, error) {
+func NewNegotiationReplyFrom(r io.Reader) (*NegotiationReply, error) {
 	bb := make([]byte, 2)
 	if _, err := io.ReadFull(r, bb); err != nil {
 		return nil, err
@@ -68,7 +67,7 @@ func NewUserPassNegotiationRequest(username []byte, password []byte) *UserPassNe
 }
 
 // WriteTo write user password negotiation request packet into server
-func (r *UserPassNegotiationRequest) WriteTo(w *net.TCPConn) error {
+func (r *UserPassNegotiationRequest) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver, r.Ulen}); err != nil {
 		return err
 	}
@@ -88,7 +87,7 @@ func (r *UserPassNegotiationRequest) WriteTo(w *net.TCPConn) error {
 }
 
 // NewUserPassNegotiationReplyFrom read user password negotiation reply packet from server
-func NewUserPassNegotiationReplyFrom(r *net.TCPConn) (*UserPassNegotiationReply, error) {
+func NewUserPassNegotiationReplyFrom(r io.Reader) (*UserPassNegotiationReply, error) {
 	bb := make([]byte, 2)
 	if _, err := io.ReadFull(r, bb); err != nil {
 		return nil, err
@@ -121,7 +120,7 @@ func NewRequest(cmd byte, atyp byte, dstaddr []byte, dstport []byte) *Request {
 }
 
 // WriteTo write request packet into server
-func (r *Request) WriteTo(w *net.TCPConn) error {
+func (r *Request) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver, r.Cmd, r.Rsv, r.Atyp}); err != nil {
 		return err
 	}
@@ -138,7 +137,7 @@ func (r *Request) WriteTo(w *net.TCPConn) error {
 }
 
 // NewReplyFrom read reply packet from server
-func NewReplyFrom(r *net.TCPConn) (*Reply, error) {
+func NewReplyFrom(r io.Reader) (*Reply, error) {
 	bb := make([]byte, 4)
 	if _, err := io.ReadFull(r, bb); err != nil {
 		return nil, err

diff --git a/server_side.go b/server_side.go
@@ -4,7 +4,6 @@ import (
 	"errors"
 	"io"
 	"log"
-	"net"
 )
 
 var (
@@ -17,7 +16,7 @@ var (
 )
 
 // NewNegotiationRequestFrom read negotiation requst packet from client
-func NewNegotiationRequestFrom(r *net.TCPConn) (*NegotiationRequest, error) {
+func NewNegotiationRequestFrom(r io.Reader) (*NegotiationRequest, error) {
 	// memory strict
 	bb := make([]byte, 2)
 	if _, err := io.ReadFull(r, bb); err != nil {
@@ -52,7 +51,7 @@ func NewNegotiationReply(method byte) *NegotiationReply {
 }
 
 // WriteTo write negotiation reply packet into client
-func (r *NegotiationReply) WriteTo(w *net.TCPConn) error {
+func (r *NegotiationReply) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver, r.Method}); err != nil {
 		return err
 	}
@@ -63,7 +62,7 @@ func (r *NegotiationReply) WriteTo(w *net.TCPConn) error {
 }
 
 // NewUserPassNegotiationRequestFrom read user password negotiation request packet from client
-func NewUserPassNegotiationRequestFrom(r *net.TCPConn) (*UserPassNegotiationRequest, error) {
+func NewUserPassNegotiationRequestFrom(r io.Reader) (*UserPassNegotiationRequest, error) {
 	bb := make([]byte, 2)
 	if _, err := io.ReadFull(r, bb); err != nil {
 		return nil, err
@@ -106,7 +105,7 @@ func NewUserPassNegotiationReply(status byte) *UserPassNegotiationReply {
 }
 
 // WriteTo write negotiation username password reply packet into client
-func (r *UserPassNegotiationReply) WriteTo(w *net.TCPConn) error {
+func (r *UserPassNegotiationReply) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver, r.Status}); err != nil {
 		return err
 	}
@@ -117,7 +116,7 @@ func (r *UserPassNegotiationReply) WriteTo(w *net.TCPConn) error {
 }
 
 // NewRequestFrom read requst packet from client
-func NewRequestFrom(r *net.TCPConn) (*Request, error) {
+func NewRequestFrom(r io.Reader) (*Request, error) {
 	bb := make([]byte, 4)
 	if _, err := io.ReadFull(r, bb); err != nil {
 		return nil, err
@@ -185,7 +184,7 @@ func NewReply(rep byte, atyp byte, bndaddr []byte, bndport []byte) *Reply {
 }
 
 // WriteTo write reply packet into client
-func (r *Reply) WriteTo(w *net.TCPConn) error {
+func (r *Reply) WriteTo(w io.Writer) error {
 	if _, err := w.Write([]byte{r.Ver, r.Rep, r.Rsv, r.Atyp}); err != nil {
 		return err
 	}