tailscale/net/proxymux/mux.go

146 lines
2.8 KiB
Go

// Copyright (c) 2021 Tailscale Inc & AUTHORS All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
// Package proxymux splits a net.Listener in two, routing SOCKS5
// connections to one and HTTP requests to the other.
//
// It allows for hosting both a SOCKS5 proxy and an HTTP proxy on the
// same listener.
package proxymux
import (
"io"
"net"
"sync"
"time"
)
// SplitSOCKSAndHTTP accepts connections on ln and passes connections
// through to either socksListener or httpListener, depending the
// first byte sent by the client.
func SplitSOCKSAndHTTP(ln net.Listener) (socksListener, httpListener net.Listener) {
sl := &listener{
addr: ln.Addr(),
c: make(chan net.Conn),
closed: make(chan struct{}),
}
hl := &listener{
addr: ln.Addr(),
c: make(chan net.Conn),
closed: make(chan struct{}),
}
go splitSOCKSAndHTTPListener(ln, sl, hl)
return sl, hl
}
func splitSOCKSAndHTTPListener(ln net.Listener, sl, hl *listener) {
for {
conn, err := ln.Accept()
if err != nil {
sl.Close()
hl.Close()
return
}
go routeConn(conn, sl, hl)
}
}
func routeConn(c net.Conn, socksListener, httpListener *listener) {
if err := c.SetReadDeadline(time.Now().Add(15 * time.Second)); err != nil {
c.Close()
return
}
var b [1]byte
if _, err := io.ReadFull(c, b[:]); err != nil {
c.Close()
return
}
if err := c.SetReadDeadline(time.Time{}); err != nil {
c.Close()
return
}
conn := &connWithOneByte{
Conn: c,
b: b[0],
}
// First byte of a SOCKS5 session is a version byte set to 5.
var ln *listener
if b[0] == 5 {
ln = socksListener
} else {
ln = httpListener
}
select {
case ln.c <- conn:
case <-ln.closed:
c.Close()
}
}
type listener struct {
addr net.Addr
c chan net.Conn
mu sync.Mutex // serializes close() on closed. It's okay to receive on closed without locking.
closed chan struct{}
}
func (ln *listener) Accept() (net.Conn, error) {
// Once closed, reliably stay closed, don't race with attempts at
// further connections.
select {
case <-ln.closed:
return nil, net.ErrClosed
default:
}
select {
case ret := <-ln.c:
return ret, nil
case <-ln.closed:
return nil, net.ErrClosed
}
}
func (ln *listener) Close() error {
ln.mu.Lock()
defer ln.mu.Unlock()
select {
case <-ln.closed:
// Already closed
default:
close(ln.closed)
}
return nil
}
func (ln *listener) Addr() net.Addr {
return ln.addr
}
// connWithOneByte is a net.Conn that returns b for the first read
// request, then forwards everything else to Conn.
type connWithOneByte struct {
net.Conn
b byte
bRead bool
}
func (c *connWithOneByte) Read(bs []byte) (int, error) {
if c.bRead {
return c.Conn.Read(bs)
}
if len(bs) == 0 {
return 0, nil
}
c.bRead = true
bs[0] = c.b
return 1, nil
}