tlsrp

tlsrp.go (6622B)

---

 package main
 
 import (
 	"context"
 	"crypto/tls"
 	"errors"
 	"flag"
 	"fmt"
 	"golang.org/x/sys/unix"
 	"io"
 	"log"
 	"net"
 	"os"
 	"os/signal"
 	"strings"
 	"sync"
 	"time"
 )
 
 // We only enforce a timeout on the handshake. After the handshake is complete,
 // the sink is responsible for timing-out clients.
 const handshakeTimeout = 10 * time.Second
 
 // softExit or hardExit is closed when exiting. During a soft exit, we
 // accept no new clients, but existing clients should finish gracefully;
 // during a hard exit, we accept no new clients, and exiting clients
 // should be forcefully disconnected. Since we have two different exit
 // modes, we don't use context.Context's to handle cancellation.
 var softExit chan struct{}
 var hardExit chan struct{}
 
 func init() {
 	softExit = make(chan struct{})
 	hardExit = make(chan struct{})
 }
 
 type socket interface {
 	io.ReadWriteCloser
 	CloseWrite() error
 }
 
 func parseSNI(sni string) (hostname, error) {
 	var hostname hostname
 
 	if sni != "" {
 		var err error
 		hostname, err = parseHostname(sni)
 		if err != nil {
 			return nil, err
 		}
 	}
 
 	return hostname, nil
 }
 
 func handshake(conn *tls.Conn) error {
 	ctx, cancel := context.WithTimeout(context.Background(), handshakeTimeout)
 	defer cancel()
 
 	handshakeErr := make(chan error, 1)
 	go func() { handshakeErr <- conn.HandshakeContext(ctx) }()
 
 	var err error
 	select {
 	case err = <-handshakeErr:
 	case <-hardExit:
 	}
 
 	return err
 }
 
 func splice(a, b socket) error {
 	a2bErr := make(chan error, 1)
 	go func() {
 		_, err := io.Copy(b, a)
 		a2bErr <- err
 	}()
 
 	b2aErr := make(chan error, 1)
 	go func() {
 		_, err := io.Copy(a, b)
 		b2aErr <- err
 	}()
 
 	// In the first two cases, we call CloseWrite (not Close) on the
 	// corresponding destination connection, so that the other copy goroutine
 	// can continue reading from it (until it hits EOF). In the hard exit
 	// case, we want to exit ASAP, so we close both ends of both connections.
 	var err error
 	select {
 	case err = <-a2bErr:
 		b.CloseWrite()
 		<-b2aErr
 	case err = <-b2aErr:
 		a.CloseWrite()
 		<-a2bErr
 	case <-hardExit:
 		a.Close()
 		b.Close()
 		<-a2bErr
 		<-b2aErr
 	}
 
 	return err
 }
 
 func proxy(client *tls.Conn) {
 	logf := func(format string, a ...any) {
 		log.Printf("client %s: %s\n", client.RemoteAddr(), fmt.Sprintf(format, a...))
 	}
 
 	logf("connected")
 	defer logf("disconnected")
 	defer client.Close()
 
 	err := handshake(client)
 	if err != nil {
 		logf("%s", err)
 		return
 	}
 
 	hostname, err := parseSNI(client.ConnectionState().ServerName)
 	if err != nil {
 		logf("rejected: %s", err)
 		return
 	}
 
 	sink, err := lookupSink(hostname)
 	if err != nil {
 		logf("rejected: %s", err)
 		return
 	}
 
 	server, err := net.Dial(sink.network, sink.address)
 	if err != nil {
 		logf("%s", err)
 		return
 	}
 	defer server.Close()
 
 	err = splice(client, server.(socket))
 	if err != nil {
 		logf("%s", err)
 		return
 	}
 }
 
 func accept(listener net.Listener) {
 	logf := func(format string, a ...any) {
 		log.Printf("source %s: %s\n", listener.Addr(), fmt.Sprintf(format, a...))
 	}
 
 	logf("accepting")
 	defer logf("closing")
 
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
 			if !errors.Is(err, net.ErrClosed) {
 				logf("%s", err)
 				os.Exit(1)
 				// XXX: Exiting here might be an over-reaction to the error.
 				// Although, keep in mind that tlsrp should be running under
 				// some service manager that should restart tlsrp if it exits.
 			}
 			return
 		}
 
 		wg.Add(1)
 		go func() {
 			proxy(conn.(*tls.Conn))
 			wg.Done()
 		}()
 	}
 }
 
 func listen(sources []string) ([]net.Listener, error) {
 	tlsConfig := &tls.Config{
 		GetCertificate: func(chi *tls.ClientHelloInfo) (*tls.Certificate, error) {
 			hostname, err := parseHostname(chi.ServerName)
 			if err != nil {
 				return nil, err
 			}
 
 			cert, err := lookupCert(hostname)
 			if err != nil {
 				return nil, err
 			}
 
 			return cert.cert, nil
 		},
 	}
 
 	listeners := make([]net.Listener, 0, len(sources))
 	for _, source := range sources {
 		fields := strings.SplitN(source, ":", 2)
 		if len(fields) != 2 {
 			return nil, fmt.Errorf("invalid source: expected colon separating network type from address")
 		}
 
 		network := fields[0]
 		address := fields[1]
 
 		switch network {
 		case "tcp", "unix":
 			listener, err := tls.Listen(network, address, tlsConfig)
 			if err != nil {
 				return nil, err
 			}
 			listeners = append(listeners, listener)
 		default:
 			return nil, fmt.Errorf("invalid source: expected network type of \"tcp\" or \"unix\"")
 		}
 	}
 
 	return listeners, nil
 }
 
 func manageConfig(cfgPath string) {
 	cfg, err := loadConfig(cfgPath)
 	if err != nil {
 		log.Printf("failed to load initial configuration: %s\n", err)
 		// Proceed with an empty configuration (i.e., with no sinks and no
 		// certs), causing every client to be rejected.
 		cfg = &config{}
 	}
 
 	sighup := make(chan os.Signal, 1)
 	signal.Notify(sighup, unix.SIGHUP)
 
 	for {
 		select {
 		case <-sighup:
 			log.Println("received SIGHUP; reloading configuration")
 			newCfg, err := loadConfig(cfgPath)
 			if err == nil {
 				cfg = newCfg
 			} else {
 				log.Printf("failed to reload configuration: %s\n", err)
 			}
 
 		case msg := <-lookupSinkChan:
 			if msg.hostname != nil {
 				for _, sink := range cfg.sinks {
 					if sink.pattern.matches(msg.hostname) {
 						msg.reply <- sink
 						break
 					}
 				}
 			} else if len(cfg.sinks) > 0 {
 				msg.reply <- cfg.sinks[0]
 			}
 			close(msg.reply)
 
 		case msg := <-lookupCertChan:
 			if msg.hostname != nil {
 				for _, cert := range cfg.certs {
 					if cert.pattern.matches(msg.hostname) {
 						msg.reply <- cert
 						break
 					}
 				}
 			} else if len(cfg.certs) > 0 {
 				msg.reply <- cfg.certs[0]
 			}
 			close(msg.reply)
 		}
 	}
 }
 
 func manageExit() {
 	sigs := make(chan os.Signal, 3)
 	signal.Notify(sigs, unix.SIGINT, unix.SIGQUIT, unix.SIGTERM)
 
 	<-sigs
 	log.Println("received SIGINT/SIGQUIT/SIGTERM; exiting softly")
 	close(softExit)
 
 	<-sigs
 	log.Println("received another SIGINT/SIGQUIT/SIGTERM; exiting harshly")
 	close(hardExit)
 }
 
 func main() {
 	flag.Usage = func() {
 		fmt.Fprintf(os.Stderr, "Usage: %s CONFIG_PATH SOURCE...", os.Args[0])
 	}
 	flag.Parse()
 	if flag.NArg() < 2 {
 		log.Fatalln("expected at least 2 arguments")
 	}
 
 	cfgPath := flag.Args()[0]
 
 	listeners, err := listen(flag.Args()[1:])
 	if err != nil {
 		log.Fatalln(err.Error())
 	}
 
 	go manageExit()
 	go manageConfig(cfgPath)
 
 	var wg sync.WaitGroup
 	defer wg.Wait()
 
 	for _, l := range listeners {
 		wg.Add(1)
 		go func (l net.Listener) {
 			accept(l)
 			wg.Done()
 		}(l)
 	}
 
 	select {
 	case <-softExit:
 	case <-hardExit:
 	}
 	for _, l := range listeners {
 		l.Close()
 	}
 }