github.com/database64128/shadowsocks-go@v1.10.2-0.20240315062903-143a773533f1/direct/stream.go (about)

     1  package direct
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/database64128/shadowsocks-go/conn"
     8  	"github.com/database64128/shadowsocks-go/socks5"
     9  	"github.com/database64128/shadowsocks-go/zerocopy"
    10  )
    11  
    12  var (
    13  	_ zerocopy.DirectReader = (*DirectStreamReadWriter)(nil)
    14  	_ zerocopy.DirectWriter = (*DirectStreamReadWriter)(nil)
    15  )
    16  
    17  // DirectStreamReadWriter implements the zerocopy ReadWriter interface and reads/writes everything
    18  // directly from/to the wrapped io.ReadWriter.
    19  type DirectStreamReadWriter struct {
    20  	rw zerocopy.DirectReadWriteCloser
    21  }
    22  
    23  // WriterInfo implements the Writer WriterInfo method.
    24  func (rw *DirectStreamReadWriter) WriterInfo() zerocopy.WriterInfo {
    25  	return zerocopy.WriterInfo{}
    26  }
    27  
    28  // WriteZeroCopy implements the Writer WriteZeroCopy method.
    29  func (rw *DirectStreamReadWriter) WriteZeroCopy(b []byte, payloadStart, payloadLen int) (payloadWritten int, err error) {
    30  	payloadWritten, err = rw.rw.Write(b[payloadStart : payloadStart+payloadLen])
    31  	return
    32  }
    33  
    34  // DirectWriter implements the DirectWriter DirectWriter method.
    35  func (rw *DirectStreamReadWriter) DirectWriter() io.Writer {
    36  	return rw.rw
    37  }
    38  
    39  // ReaderInfo implements the Reader ReaderInfo method.
    40  func (rw *DirectStreamReadWriter) ReaderInfo() zerocopy.ReaderInfo {
    41  	return zerocopy.ReaderInfo{}
    42  }
    43  
    44  // ReadZeroCopy implements the Reader ReadZeroCopy method.
    45  func (rw *DirectStreamReadWriter) ReadZeroCopy(b []byte, payloadBufStart, payloadBufLen int) (payloadLen int, err error) {
    46  	payloadLen, err = rw.rw.Read(b[payloadBufStart : payloadBufStart+payloadBufLen])
    47  	return
    48  }
    49  
    50  // DirectReader implements the DirectReader DirectReader method.
    51  func (rw *DirectStreamReadWriter) DirectReader() io.Reader {
    52  	return rw.rw
    53  }
    54  
    55  // CloseWrite implements the ReadWriter CloseWrite method.
    56  func (rw *DirectStreamReadWriter) CloseWrite() error {
    57  	return rw.rw.CloseWrite()
    58  }
    59  
    60  // CloseRead implements the ReadWriter CloseRead method.
    61  func (rw *DirectStreamReadWriter) CloseRead() error {
    62  	return rw.rw.CloseRead()
    63  }
    64  
    65  // Close implements the ReadWriter Close method.
    66  func (rw *DirectStreamReadWriter) Close() error {
    67  	return rw.rw.Close()
    68  }
    69  
    70  // NewDirectStreamReadWriter returns a ReadWriter that passes all reads and writes directly to the underlying stream.
    71  func NewDirectStreamReadWriter(rw zerocopy.DirectReadWriteCloser) *DirectStreamReadWriter {
    72  	return &DirectStreamReadWriter{
    73  		rw: rw,
    74  	}
    75  }
    76  
    77  // NewShadowsocksNoneStreamClientReadWriter creates a ReadWriter that acts as a Shadowsocks none method client.
    78  func NewShadowsocksNoneStreamClientReadWriter(ctx context.Context, rwo zerocopy.DirectReadWriteCloserOpener, targetAddr conn.Addr, payload []byte) (*DirectStreamReadWriter, zerocopy.DirectReadWriteCloser, error) {
    79  	targetAddrLen := socks5.LengthOfAddrFromConnAddr(targetAddr)
    80  	writeBuf := make([]byte, targetAddrLen+len(payload))
    81  	socks5.WriteAddrFromConnAddr(writeBuf, targetAddr)
    82  	copy(writeBuf[targetAddrLen:], payload)
    83  	rawRW, err := rwo.Open(ctx, writeBuf)
    84  	if err != nil {
    85  		return nil, nil, err
    86  	}
    87  	return &DirectStreamReadWriter{rw: rawRW}, rawRW, nil
    88  }
    89  
    90  // NewShadowsocksNoneStreamServerReadWriter creates a ReadWriter that acts as a Shadowsocks none method server.
    91  func NewShadowsocksNoneStreamServerReadWriter(rw zerocopy.DirectReadWriteCloser) (*DirectStreamReadWriter, conn.Addr, error) {
    92  	addr, err := socks5.ConnAddrFromReader(rw)
    93  	if err != nil {
    94  		return nil, addr, err
    95  	}
    96  	return &DirectStreamReadWriter{rw: rw}, addr, nil
    97  }
    98  
    99  // NewSocks5StreamClientReadWriter writes a SOCKS5 CONNECT request to rw and wraps rw into a ReadWriter ready for use.
   100  func NewSocks5StreamClientReadWriter(rw zerocopy.DirectReadWriteCloser, targetAddr conn.Addr) (*DirectStreamReadWriter, error) {
   101  	if err := socks5.ClientConnect(rw, targetAddr); err != nil {
   102  		return nil, err
   103  	}
   104  	return &DirectStreamReadWriter{rw: rw}, nil
   105  }
   106  
   107  // NewSocks5StreamServerReadWriter handles a SOCKS5 request from rw and wraps rw into a ReadWriter ready for use.
   108  // conn must be provided when UDP is enabled.
   109  func NewSocks5StreamServerReadWriter(rw zerocopy.DirectReadWriteCloser, enableTCP, enableUDP bool) (dsrw *DirectStreamReadWriter, addr conn.Addr, err error) {
   110  	addr, err = socks5.ServerAccept(rw, enableTCP, enableUDP)
   111  	if err == nil {
   112  		dsrw = &DirectStreamReadWriter{
   113  			rw: rw,
   114  		}
   115  	}
   116  	return
   117  }