dumbproxy.go

package main

import (
	"flag"
	"fmt"
	"io"
	"net"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"
)

var (
	logfile = flag.String("l", "", "log filename")
	host    = flag.String("h", "", "target hostname")
	port    = flag.Uint("p", 22, "target port")
	logf    = os.Stderr
)

func fatal(format string, args ...interface{}) {
	os.Stderr.WriteString(fmt.Sprintf(format+"\n", args...))
	os.Exit(1)
}

func log(format string, args ...interface{}) {
	prefix := time.Now().Format(time.StampMilli) + " > "
	logf.WriteString(prefix + fmt.Sprintf(format, args...) + "\n")
}

func main() {
	flag.Parse()
	if *host == "" || len(flag.Args()) > 0 || *port >= (1<<16) {
		flag.PrintDefaults()
		os.Exit(2)
	}

	if *logfile != "" {
		f, err := os.OpenFile(*logfile, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0644)
		if err != nil {
			fatal("ERROR opening log file: %s", err)
		}
		logf = f
	}

	start_ts := time.Now()

	addr := net.JoinHostPort(*host, strconv.Itoa(int(*port)))
	conn, err := net.Dial("tcp", addr)
	if err != nil {
		fatal("ERROR connecting to target: %s", err)
	}

	log("connected to %s in %fs", addr, time.Since(start_ts).Seconds())

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	closed_send := make(chan int)
	closed_recv := make(chan int)

	bytes_sent := 0
	bytes_recv := 0

	go func() {
		for {
			var buf [4096]byte
			r, err := os.Stdin.Read(buf[:])
			if err != nil {
				if err != io.EOF {
					log("ERROR reading stdin: %s", err)
				}
				conn.(*net.TCPConn).CloseWrite()
				break
			}

			w, err := conn.Write(buf[:r])
			if err != nil {
				log("ERROR in send: %s", err)
				break
			}
			bytes_sent += w

			if w != r {
				log("ERROR small send %d != %d", w, r)
				break
			}
		}
		close(closed_send)
	}()

	go func() {
		for {
			var buf [4096]byte
			r, err := conn.Read(buf[:])
			if err != nil {
				if err != io.EOF {
					log("ERROR in recv: %s", err)
				}
				os.Stdout.Close()
				break
			}
			bytes_recv += r

			w, err := os.Stdout.Write(buf[:r])
			if err != nil {
				log("ERROR writing stdout: %s", err)
				break
			}

			if w != r {
				log("ERROR small write %d != %d", w, r)
				break
			}
		}
		close(closed_recv)
	}()

	select {
	case <-sigChan:
		log("received exit signal")
		sigChan = nil
		break
	case <-closed_send:
		log("stdin closed")
		closed_send = nil
		break
	case <-closed_recv:
		log("connection closed")
		closed_recv = nil
		break
	}
	//conn.Close()
	log("finished connection to %s sent=%d recv=%d duration=%fs",
		addr, bytes_sent, bytes_recv, time.Since(start_ts).Seconds())
}