Improve connection handling

main.go

@@ -13,68 +13,94 @@ var connect string
 
 func main() {
 	flag.IntVar(&key, "key", 170, "key to xor the data")
-	flag.StringVar(&listen, "listen", ":8000", "listen on ip and port")
-	flag.StringVar(&connect, "connect", "", "forward to ip and port")
+	flag.StringVar(&listen, "listen", ":8081", "listen on ip and port")
+	flag.StringVar(&connect, "connect", ":8080", "forward to ip and port")
 	flag.Parse()
 
 	if key < 0 || key > 255 {
-        flag.PrintDefaults()
+		flag.PrintDefaults()
 		log.Fatal("key is not one byte")
 	}
 
-    if len(connect) < 3 {
-        flag.PrintDefaults()
-        log.Fatal("no connection specified")
-    }
+	// check and parse address
+	conn, err := net.ResolveTCPAddr("tcp", connect)
+	if err != nil {
+		flag.PrintDefaults()
+		log.Fatal(err)
+	}
 
+	// listen on address
 	ln, err := net.Listen("tcp", listen)
 	if err != nil {
 		flag.PrintDefaults()
 		log.Fatal(err)
 	}
-	log.Print("listening on " + listen)
-	log.Print("will connect to " + connect)
 
-	for {
+	log.Printf("listening on %v", ln.Addr())
+	log.Printf("will connect to %v", conn)
+
+	for i := 0; ; i++ {
+		// accept new connection
 		c, err := ln.Accept()
 		if err != nil {
-			log.Print(err)
-            continue
-		}	
-        log.Print("Connection from "+c.RemoteAddr().String())
+			log.Fatal(err)
+		}
 
-		cn, err := net.Dial("tcp", connect)
+		log.Printf("connection %v from %v", i, c.RemoteAddr())
+
+		cn, err := net.DialTCP("tcp", nil, conn)
 		if err != nil {
+			c.Close()
 			log.Print(err)
-            continue
+			continue
 		}
 
-		go pipe(c, cn, byte(key))
-		go pipe(cn, c, byte(key))
+		go pipe(c, cn, byte(key), i)
+		go pipe(cn, c, byte(key), i)
 	}
 }
 
-func pipe(w io.WriteCloser, r io.ReadCloser, key byte) {
-	buff := make([]byte, 65535)
+func pipe(w io.WriteCloser, r io.ReadCloser, key byte, count int) {
+	n, err := copyBufferXor(w, r, key)
+
+	r.Close()
+	w.Close()
+
+	log.Printf("connection %v closed, %v bytes", count, n)
+
+	opError, ok := err.(*net.OpError)
+	if err != nil && (!ok || opError.Op != "readfrom") {
+		log.Printf("warning! %v", err)
+	}
+}
+
+func copyBufferXor(dst io.Writer, src io.Reader, key byte) (written int64, err error) {
+	buf := make([]byte, 32*1024)
 	for {
-		n, rerr := r.Read(buff)
-		for i := 0; i < n; i++ {
-			buff[i] = buff[i] ^ key
+		nr, er := src.Read(buf)
+		for i := 0; i < nr; i++ {
+			buf[i] = buf[i] ^ key
 		}
-		wn, werr := w.Write(buff[:n])
-        if n != wn {
-            log.Print("mismatch")
-        }
-
-		if werr != nil {
-			r.Close()
-			log.Print(werr)
-            return
+		if nr > 0 {
+			nw, ew := dst.Write(buf[0:nr])
+			if nw > 0 {
+				written += int64(nw)
+			}
+			if ew != nil {
+				err = ew
+				break
+			}
+			if nr != nw {
+				err = io.ErrShortWrite
+				break
+			}
 		}
-		if rerr != nil {
-			w.Close()
-			log.Print(rerr)
-            return
+		if er != nil {
+			if er != io.EOF {
+				err = er
+			}
+			break
 		}
 	}
+	return written, err
 }