mitm/main.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"time"
)

type fproxy struct{}

type peek struct {
	buff  *bytes.Buffer
	buff2 *bytes.Buffer
}

func newPeek() peek {
	return peek{
		buff:  bytes.NewBuffer(nil),
		buff2: bytes.NewBuffer(nil),
	}
}

func (peek peek) Read(b []byte) (int, error) {
	n, err := peek.buff.Read(b)
	peek.buff2.Write(b[:n])
	log.Printf("read (%d): %s", n, b[:n])
	return n, err
}

func (peek peek) Write(b []byte) (int, error) {
	n, err := peek.buff.Write(b)
	log.Printf("wrote %d): %s", n, b[:n])
	return n, err
}

func (peek peek) Close() error {
	//log.Printf("Close: %s", peek.buff2.Bytes())
	return nil
}

func main() {
	port := flag.String("p", ":9005", "addr to listen on")
	flag.Parse()
	log.Printf("listening on %s", *port)
	if err := http.ListenAndServe(*port, fproxy{}); err != nil {
		panic(err)
	}
}

func (fproxy fproxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if err := fproxy.serveHTTP(w, r); err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}
}

func (fproxy fproxy) serveHTTP(w http.ResponseWriter, r *http.Request) error {
	if r.Method == http.MethodConnect {
		return fproxy.passthrough(w, r)
	}
	return fproxy.proxy(w, r)
}

func (fproxy fproxy) proxy(w http.ResponseWriter, r *http.Request) error {
	if r.URL.Scheme == "" {
		r.URL.Scheme = "http"
	}
	b, err := ioutil.ReadAll(r.Body)
	if err != nil {
		return err
	}
	r.Body = ioutil.NopCloser(bytes.NewReader(b))
	resp, err := (&http.Transport{}).RoundTrip(r)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	for k, v := range resp.Header {
		for i := range v {
			w.Header().Add(k, v[i])
		}
	}
	w.WriteHeader(resp.StatusCode)
	peek := bytes.NewBuffer(nil)
	forward := func(b []byte) {
		peek.Write(b)
		w.Write(b)
	}
	buff := make([]byte, 1024)
	var n int
	for n, err = resp.Body.Read(buff); err == nil; n, err = resp.Body.Read(buff) {
		forward(buff[:n])
	}
	forward(buff[:n])
	if err == io.EOF {
		err = nil
	}
	log.Printf("%s %s: %+v: %s => %d: %+v: %s", r.Method, r.URL.String(), r.Header, b, resp.StatusCode, resp.Header, peek.Bytes())
	//_, err = io.Copy(w, resp.Body)
	return err
}

func (fproxy fproxy) passthrough(w http.ResponseWriter, r *http.Request) error {
	log.Printf("passthrough")
	defer log.Printf("/passthrough")

	dest_conn, err := net.DialTimeout("tcp", r.Host, 30*time.Second)
	if err != nil {
		return err
	}
	w.WriteHeader(http.StatusOK)
	hijacker, ok := w.(http.Hijacker)
	if !ok {
		return errors.New("Hijacking not supported")
	}
	client_conn, _, err := hijacker.Hijack()
	if err != nil {
		return err
	}

	transfer := func(destination io.WriteCloser, source io.ReadCloser) {
		defer destination.Close()
		defer source.Close()
		io.Copy(destination, source)
		return
		/*
			b := make([]byte, 1024)
			for n, err := source.Read(b); err != io.EOF; n, err = source.Read(b) {
				log.Printf("read %v/%v", n, err)
				if err != nil {
					xferErr = err
					return
				}
				for n > 0 {
					m, err := destination.Write(b[:n])
					if err != nil {
						xferErr = err
						return
					}
					n -= m
				}
			}
		*/
	}

	go transfer(dest_conn, client_conn)
	go transfer(client_conn, dest_conn)

	return nil
}