tcp relay with rc4;;

main.c

package main

import (
	"bytes"
	"crypto/rand"
	"crypto/rc4"
	"crypto/sha256"
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"os/signal"
	"time"
)

func errorPanic(err error, format string, a ...interface{}) {
	if err != nil {
		fmt.Printf("Panic : "+format, a)
		panic(err)
	}
}

/*
{
    "Relays":[
        {
            "Ip":"",
            "Port":0,
            "Remote":"",
            "Key":""
        },
        {
            "Ip":"",
            "Port":0,
            "Remote":"",
            "Key":""
        }
    ],
    "Client":false
}
 */

type relay struct {
	Ip     string `json:ip`
	Port   int    `json:port`
	Remote string `json:remote`
	Key    string `json:key`
}

type config struct {
	Relays []relay `json:relays`
	Client bool    `json:client`
}

func (c *config) load(f string) {
	jsonBytes, err := ioutil.ReadFile(f)
	errorPanic(err, "open file %s error %+v \n", f, err)
	err = json.Unmarshal(jsonBytes, c)
	errorPanic(err, "Unmarshal file %s error %+v \n", f, err)
}

type tcpRelay struct {
	listenConn *net.TCPListener
	remote     string
	key        []byte
	client     bool
}

const nonceSize int = 8

func getNonce(nonce []byte) error {
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		fmt.Println("rand.Reader error: ", err)
		return err
	}
	return nil
}

func newTcpRelay(c* relay, client bool) *tcpRelay {
	ret := &tcpRelay{}
	ret.client = client
	keyHash := sha256.Sum256([]byte(c.Key))
	ret.key = keyHash[:]
	ret.remote = c.Remote
	var err error
	ret.listenConn, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(c.Ip), Port: c.Port})
	errorPanic(err, "open ListenTCP error +%v", err)
	go ret.handleListen()
	return ret
}

func (r *tcpRelay) handleListen() {
	for {
		c, err := r.listenConn.Accept()
		errorPanic(err, "accept error: ", err)
		if r.client {
			go r.handleConnClient(c)
		} else {
			go r.handleConnServer(c)
		}

	}
}

func (r *tcpRelay) handleConnClient(client net.Conn) {
	remote, err := net.DialTimeout("tcp", r.remote, 2 * time.Second)
	if err != nil {
		fmt.Println("DialTimeout error:", err)
		client.Close()
		return
	}
	tcp := &tcpConn{}
	tcp.cipherConn = remote
	tcp.plainConn = client
	tcp.key = r.key
	go tcp.handleEncrypt()
	go tcp.handleDecrypt()
}

func (r *tcpRelay) handleConnServer(client net.Conn) {
	remote, err := net.DialTimeout("tcp", r.remote, 2 * time.Second)
	if err != nil {
		fmt.Println("DialTimeout error:", err)
		client.Close()
		return
	}
	tcp := &tcpConn{}
	tcp.cipherConn = client
	tcp.plainConn = remote
	tcp.key = r.key
	go tcp.handleEncrypt()
	go tcp.handleDecrypt()
}

func (r *tcpRelay) stop() {
	r.listenConn.Close()
}

func tcpWrite(tcp net.Conn, buf []byte) error {
	n, err := tcp.Write(buf)
	if err != nil {
		return err
	}
	if n != len(buf) {
		return fmt.Errorf("tcp.Write n != len")
	}
	return nil
}

func tcpRead(tcp net.Conn, buf []byte, need int) error {
	pos := 0
	for {
		n, err := tcp.Read(buf[pos:need])
		if err != nil {
			return fmt.Errorf("tcp.cipherConn.Read(header) error: %+v", err)
		}
		if n == 0 {
			return fmt.Errorf("tcp.cipherConn.Read(header) n == 0: %+v", err)
		}
		pos += n
		if pos == need {
			break
		}
	}
	return nil
}

type tcpConn struct {
	cipherConn net.Conn
	plainConn  net.Conn
	key        []byte
}

func (tcp *tcpConn) getRc4(nonce []byte) (*rc4.Cipher, error) {
	sha := sha256.New()
	sha.Write(nonce)
	sha.Write(tcp.key)
	return rc4.NewCipher(sha.Sum(nil))
}

func (tcp *tcpConn) handleEncrypt() {
	defer tcp.plainConn.Close()
	defer tcp.cipherConn.Close()
	buffer := make([]byte, 1500)
	data := buffer[2:]
	length := buffer[:2]
	header := make([]byte, nonceSize + sha256.Size)
	copy(header[nonceSize:], tcp.key)
	if err := getNonce(header[:nonceSize]); err != nil {
		fmt.Println("getNonce error:", err)
		return
	}
	encryptRc4, err := tcp.getRc4(header[:nonceSize])
	if err != nil {
		fmt.Println("getRc4 error:", err)
		return
	}
	encryptRc4.XORKeyStream(header[nonceSize:], header[nonceSize:])
	if err := tcpWrite(tcp.cipherConn, header); err != nil {
		fmt.Println("tcpWrite header error:", err)
		return
	}
	for {
		n, err := tcp.plainConn.Read(data)
		if err != nil {
			fmt.Println("tcp.plainConn.Read(data) error:", err)
			return
		}
		if n == 0 {
			fmt.Println("tcp.plainConn.Read(data) n == 0:", err)
			return
		}
		binary.BigEndian.PutUint16(length, uint16(n))
		toSend := buffer[:n+2]
		encryptRc4.XORKeyStream(toSend, toSend)
		if err := tcpWrite(tcp.cipherConn, toSend); err != nil {
			fmt.Println("tcpWrite tcp.cipherConn toSend error:", err)
			return
		}
	}
}

func (tcp *tcpConn) handleDecrypt() {
	defer tcp.plainConn.Close()
	defer tcp.cipherConn.Close()
	buffer := make([]byte, 1500)
	data := buffer[2:]
	length := buffer[:2]
	header := make([]byte, nonceSize + sha256.Size)
	var decryptRc4 *rc4.Cipher = nil
	if err := tcpRead(tcp.cipherConn, header, nonceSize + sha256.Size); err != nil {
		fmt.Println("tcpRead header error:", err)
		return
	}
	if r, err := tcp.getRc4(header[:nonceSize]); err != nil {
		fmt.Println("tcp.getRc4(header[:nonceSize]) error:", err)
		return
	} else {
		decryptRc4 = r
		decryptRc4.XORKeyStream(header[nonceSize:], header[nonceSize:])
		if bytes.Equal(header[nonceSize:], tcp.key) == false {
			fmt.Println("bytes.Equal == false", err)
			return
		}
	}
	for {
		if err := tcpRead(tcp.cipherConn, length, 2); err != nil {
			fmt.Println("tcpRead length error:", err)
			return
		}
		decryptRc4.XORKeyStream(length, length)
		need := int(binary.BigEndian.Uint16(length))
		if need > 1498 {
			fmt.Println("length > 1498")
			return
		}
		if err := tcpRead(tcp.cipherConn, data, need); err != nil {
			fmt.Println("tcpRead data error:", err)
			return
		}
		toSend := data[:need]
		decryptRc4.XORKeyStream(toSend, toSend)
		if err := tcpWrite(tcp.plainConn, toSend); err != nil {
			fmt.Println("tcpWrite tcp.plainConn toSend error:", err)
			return
		}
	}
}


func main() {
	c := &config{}
	if len(os.Args) > 1 {
		c.load(os.Args[1])
	} else {
		c.load("./conf.json")
	}
	relays := make([]*tcpRelay, len(c.Relays))
	for pos := range c.Relays {
		relays[pos] = newTcpRelay(&c.Relays[pos], c.Client)
	}
	signalChan := make(chan os.Signal, 2)
	signal.Notify(signalChan, os.Interrupt)
	<-signalChan
	for _, r := range relays {
		r.stop()
	}
}