email/cmd/ximap/proxy/proxy.go
2016-10-09 21:21:25 -07:00

84 lines
2.0 KiB
Go

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++
}
}