package main

import (
	"bufio"
	"bytes"
	"flag"
	"fmt"
	"net"

	"github.com/golang/glog"
)

var (
	saddr = flag.String("saddr", ":11143", "address to listen on")
	caddr = flag.String("caddr", "localhost:10143", "remote address to connect to")
)

// scanLines reimplements bufio.ScanLines without eating '\r' in the input.
func scanLines(data []byte, atEOF bool) (advance int, token []byte, err error) {
	if atEOF && len(data) == 0 {
		return 0, nil, nil
	}
	if i := bytes.IndexByte(data, '\n'); i >= 0 {
		// We have a full newline-terminated line.
		return i + 1, data[0:i], nil
	}
	// If we're at EOF, we have a final, non-terminated line. Return it.
	if atEOF {
		return len(data), data, nil
	}
	// Request more data.
	return 0, nil, nil
}

func prefixCopy(prefix string, src, dst net.Conn) {
	defer src.Close()
	defer dst.Close()

	glog.Infof("%s proxying %s -> %s", prefix, src.LocalAddr(), dst.RemoteAddr())
	scanner := bufio.NewScanner(src)
	scanner.Split(scanLines)
	for scanner.Scan() {
		glog.Infoln(prefix, scanner.Text()) // Println will add back the final '\n'
		if _, err := fmt.Fprintln(dst, scanner.Text()); err != nil {
			glog.Errorf("%s error writing to dst: %v", prefix, err)
			return
		}
	}
	if err := scanner.Err(); err != nil {
		glog.Errorf("%s error reading from src: %v", prefix, err)
		return
	}
}

func proxy(id int, client net.Conn, remoteAddr string) {
	remote, err := net.Dial("tcp", remoteAddr)
	if err != nil {
		glog.Errorf("Failed to dial %s: %v", remoteAddr, err)
		return
	}
	go prefixCopy(fmt.Sprintf("[%d] S:", id), remote, client)
	go prefixCopy(fmt.Sprintf("[%d] C:", id), client, remote)
}

func main() {
	flag.Parse()
	defer glog.Flush()

	glog.Infof("Proxying %s -> %s", *saddr, *caddr)
	ln, err := net.Listen("tcp", *saddr)
	if err != nil {
		glog.Exitf("Failed to listen on %s: %v", *saddr, err)
	}
	var id int
	for {
		conn, err := ln.Accept()
		if err != nil {
			glog.Exitf("Failed to accept on %s: %v", *saddr, err)
		}
		proxy(id, conn, *caddr)
		id++
	}
}