github.com/ipfans/trojan-go@v0.11.0/tunnel/mux/conn.go (about)

     1  package mux
     2  
     3  import (
     4  	"io"
     5  	"math/rand"
     6  
     7  	"github.com/ipfans/trojan-go/log"
     8  	"github.com/ipfans/trojan-go/tunnel"
     9  )
    10  
    11  type stickyConn struct {
    12  	tunnel.Conn
    13  	synQueue chan []byte
    14  	finQueue chan []byte
    15  }
    16  
    17  func (c *stickyConn) stickToPayload(p []byte) []byte {
    18  	buf := make([]byte, 0, len(p)+16)
    19  	for {
    20  		select {
    21  		case header := <-c.synQueue:
    22  			buf = append(buf, header...)
    23  		default:
    24  			goto stick1
    25  		}
    26  	}
    27  stick1:
    28  	buf = append(buf, p...)
    29  	for {
    30  		select {
    31  		case header := <-c.finQueue:
    32  			buf = append(buf, header...)
    33  		default:
    34  			goto stick2
    35  		}
    36  	}
    37  stick2:
    38  	return buf
    39  }
    40  
    41  func (c *stickyConn) Close() error {
    42  	const maxPaddingLength = 512
    43  	padding := [maxPaddingLength + 8]byte{'A', 'B', 'C', 'D', 'E', 'F'} // for debugging
    44  	buf := c.stickToPayload(nil)
    45  	c.Write(append(buf, padding[:rand.Intn(maxPaddingLength)]...))
    46  	return c.Conn.Close()
    47  }
    48  
    49  func (c *stickyConn) Write(p []byte) (int, error) {
    50  	if len(p) == 8 {
    51  		if p[0] == 1 || p[0] == 2 { // smux 8 bytes header
    52  			switch p[1] {
    53  			// THE CONTENT OF THE BUFFER MIGHT CHANGE
    54  			// NEVER STORE THE POINTER TO HEADER, COPY THE HEADER INSTEAD
    55  			case 0:
    56  				// cmdSYN
    57  				header := make([]byte, 8)
    58  				copy(header, p)
    59  				c.synQueue <- header
    60  				return 8, nil
    61  			case 1:
    62  				// cmdFIN
    63  				header := make([]byte, 8)
    64  				copy(header, p)
    65  				c.finQueue <- header
    66  				return 8, nil
    67  			}
    68  		} else {
    69  			log.Debug("other 8 bytes header")
    70  		}
    71  	}
    72  	_, err := c.Conn.Write(c.stickToPayload(p))
    73  	return len(p), err
    74  }
    75  
    76  func newStickyConn(conn tunnel.Conn) *stickyConn {
    77  	return &stickyConn{
    78  		Conn:     conn,
    79  		synQueue: make(chan []byte, 128),
    80  		finQueue: make(chan []byte, 128),
    81  	}
    82  }
    83  
    84  type Conn struct {
    85  	rwc io.ReadWriteCloser
    86  	tunnel.Conn
    87  }
    88  
    89  func (c *Conn) Read(p []byte) (int, error) {
    90  	return c.rwc.Read(p)
    91  }
    92  
    93  func (c *Conn) Write(p []byte) (int, error) {
    94  	return c.rwc.Write(p)
    95  }
    96  
    97  func (c *Conn) Close() error {
    98  	return c.rwc.Close()
    99  }