tunnel.go

package main

import (
        "crypto/rand"
        "errors"
        "flag"
        "fmt"
        "io"
        "net"

        "golang.org/x/crypto/chacha20"
)

var globalk []byte = []byte("xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx")

func main() {
        localAddr := flag.String("la", "localhost:2000", "")
        remoteAddr := flag.String("ra", "localhost:2001", "")
        role := flag.String("role", "", "A or B")
        flag.Parse()

        listener, err := net.Listen("tcp", *localAddr)
        if err != nil {
                fmt.Printf("Listen failed: %v\n", err)
                return
        }
        for {
                conn, err := listener.Accept()
                if err != nil {
                        fmt.Printf("Accept failed: %v", err)
                        continue
                }
                dest, err := net.Dial("tcp", *remoteAddr)
                if err != nil {
                        panic(err)
                }

                var cdest, cconn CipherStream
                switch *role {
                case "A":
                        cdest, err = NewChacha20CipherStream(globalk, dest)
                        if err != nil {
                                panic(err)
                        }
                        cconn = conn
                case "B":
                        cconn, err = NewChacha20CipherStream(globalk, conn)
                        if err != nil {
                                panic(err)
                        }
                        cdest = dest
                }
                go forward(cconn, cdest)
        }

}

func forward(client, dest CipherStream) {
        f := func(client, dest CipherStream) {
                defer client.Close()
                defer dest.Close()
                _, _ = io.Copy(client, dest)
        }

        go f(client, dest)
        go f(dest, client)
}

type Chacha20CipherStream struct {
        key     []byte // key is shared by both ends, and should be 32B
        encoder *chacha20.Cipher
        decoder *chacha20.Cipher
        conn    net.Conn
}

type CipherStream interface {
        Read([]byte) (int, error)
        Write([]byte) (int, error)
        Close() error
}

func NewChacha20CipherStream(key []byte, conn net.Conn) (*Chacha20CipherStream, error) {
        s := &Chacha20CipherStream{key, nil, nil, conn}
        nonce := make([]byte, chacha20.NonceSizeX)
        var err error
        if _, err = rand.Read(nonce); err != nil {
                return s, err
        }
        s.encoder, err = chacha20.NewUnauthenticatedCipher(s.key, nonce)
        if err != nil {
                return s, err
        }

        if n, err := s.conn.Write(nonce); err != nil || n != len(nonce) {
                return s, errors.New("write nonce failed: " + err.Error())
        }
        return s, nil

}

/* Overwrite net.Conn: Read, Write, Close */
func (s *Chacha20CipherStream) Read(b []byte) (int, error) {
        if s.decoder == nil {
                nonce := make([]byte, chacha20.NonceSizeX)
                if n, err := io.ReadAtLeast(s.conn, nonce, len(nonce)); n != len(nonce) || err != nil {
                        return n, errors.New("can't read nonce from stream: " + err.Error())
                }
                decoder, err := chacha20.NewUnauthenticatedCipher(s.key, nonce)
                if err != nil {
                        return 0, err
                }
                s.decoder = decoder
        }
        n, err := s.conn.Read(b)
        if n == 0 || err != nil {
                return n, err
        }
        bn := b[:n]
        s.decoder.XORKeyStream(bn, bn)
        copy(b, bn)
        return n, nil

}

func (s *Chacha20CipherStream) Write(b []byte) (int, error) {
        dst := make([]byte, len(b))
        s.encoder.XORKeyStream(dst, b)
        return s.conn.Write(dst)
}
func (s *Chacha20CipherStream) Close() error {
        return s.conn.Close()

}