Merge pull request #169 from SenseUnit/pipeconn

Pipeconn

Author: Snawoot

dialer/pipewrap.go

@@ -0,0 +1,131 @@
+package dialer
+
+import (
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/hashicorp/go-multierror"
+)
+
+type ReadPipe interface {
+	io.Reader
+	io.WriterTo
+	io.Closer
+	Fd() uintptr
+	SetReadDeadline(t time.Time) error
+}
+
+type WritePipe interface {
+	io.Writer
+	io.ReaderFrom
+	io.Closer
+	Fd() uintptr
+	SetWriteDeadline(t time.Time) error
+}
+
+type PipeAddr struct {
+	rfd uintptr
+	wfd uintptr
+}
+
+func (_ PipeAddr) Network() string {
+	return "pipe"
+}
+
+func (a PipeAddr) String() string {
+	return fmt.Sprintf("<read fd: %d, write rd: %d>", a.rfd, a.wfd)
+}
+
+type PipeConn struct {
+	r  ReadPipe
+	w  WritePipe
+	rc sync.Once
+	wc sync.Once
+}
+
+func NewPipeConn(r ReadPipe, w WritePipe) *PipeConn {
+	return &PipeConn{
+		r: r,
+		w: w,
+	}
+}
+
+func (c *PipeConn) Read(p []byte) (n int, err error) {
+	return c.r.Read(p)
+}
+
+func (c *PipeConn) Write(p []byte) (n int, err error) {
+	return c.w.Write(p)
+}
+
+func (c *PipeConn) Close() error {
+	var err error
+	if closeErr := c.CloseWrite(); closeErr != nil {
+		err = multierror.Append(err, closeErr)
+	}
+	if closeErr := c.CloseRead(); closeErr != nil {
+		err = multierror.Append(err, closeErr)
+	}
+	return err
+}
+
+func (c *PipeConn) CloseWrite() error {
+	var err error
+	c.wc.Do(func() {
+		err = c.w.Close()
+	})
+	return err
+}
+
+func (c *PipeConn) CloseRead() error {
+	var err error
+	c.wc.Do(func() {
+		err = c.r.Close()
+	})
+	return err
+}
+
+func (c *PipeConn) LocalAddr() net.Addr {
+	return PipeAddr{
+		rfd: c.r.Fd(),
+		wfd: c.w.Fd(),
+	}
+}
+
+func (c *PipeConn) RemoteAddr() net.Addr {
+	return c.LocalAddr()
+}
+
+func (c *PipeConn) SetReadDeadline(t time.Time) error {
+	return c.r.SetReadDeadline(t)
+}
+
+func (c *PipeConn) SetWriteDeadline(t time.Time) error {
+	return c.w.SetWriteDeadline(t)
+}
+
+func (c *PipeConn) SetDeadline(t time.Time) error {
+	var err error
+	if cErr := c.SetReadDeadline(t); err != nil {
+		err = multierror.Append(err, cErr)
+	}
+	if cErr := c.SetWriteDeadline(t); err != nil {
+		err = multierror.Append(err, cErr)
+	}
+	return err
+}
+
+func (c *PipeConn) ReadFrom(r io.Reader) (n int64, err error) {
+	return c.w.ReadFrom(r)
+}
+
+func (c *PipeConn) WriteTo(w io.Writer) (n int64, err error) {
+	return c.r.WriteTo(w)
+}
+
+var _ net.Conn = new(PipeConn)
+var _ io.ReaderFrom = new(PipeConn)
+var _ io.WriterTo = new(PipeConn)

handler/stdio.go

@@ -3,23 +3,24 @@ package handler
 import (
 	"context"
 	"fmt"
-	"io"
 	"net"
 	"sync"
 
 	clog "github.com/SenseUnit/dumbproxy/log"
+
+	"github.com/SenseUnit/dumbproxy/dialer"
 )
 
-func StdIOHandler(dialer HandlerDialer, logger *clog.CondLogger, forward ForwardFunc) func(ctx context.Context, reader io.Reader, writer io.Writer, dstAddress string) error {
-	return func(ctx context.Context, reader io.Reader, writer io.Writer, dstAddress string) error {
+func StdIOHandler(d HandlerDialer, logger *clog.CondLogger, forward ForwardFunc) func(context.Context, dialer.ReadPipe, dialer.WritePipe, string) error {
+	return func(ctx context.Context, reader dialer.ReadPipe, writer dialer.WritePipe, dstAddress string) error {
 		logger.Debug("Request: %v => %v %q %v %v %v", "<stdio>", "<stdio>", "", "STDIO", "CONNECT", dstAddress)
-		target, err := dialer.DialContext(ctx, "tcp", dstAddress)
+		target, err := d.DialContext(ctx, "tcp", dstAddress)
 		if err != nil {
 			return fmt.Errorf("connect to %q failed: %w", dstAddress, err)
 		}
 		defer target.Close()
 
-		return forward(ctx, "", wrapSOCKS(reader, writer), target)
+		return forward(ctx, "", dialer.NewPipeConn(reader, writer), target)
 	}
 }