github.com/database64128/shadowsocks-go@v1.7.0/direct/stream.go (about)

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