port forwarding with yamux (Go)

ytun.go

package main

import (
	"flag"
	"github.com/hashicorp/yamux"
	"io"
	"log"
	"net"
	"sync"
)

var local_addr = flag.String("l", "", "local addr")
var remote_addr = flag.String("r", "", "remote addr")
var forward_addr = flag.String("f", "", "forward addr")

func main() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	flag.Parse()

	if len(*forward_addr) > 0 {
		run_server()
	} else {
		run_client()
	}

}

func loop_accept(f func(conn net.Conn)) {
	ln, err := net.Listen("tcp", *local_addr)
	if err != nil {
		log.Fatal(err)
	}

	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Fatal(err)
		}
		defer conn.Close()

		f(conn)
	}
}

func concat_conn(src, dst net.Conn) {
	defer src.Close()
	defer dst.Close()

	var wg sync.WaitGroup

	cp := func(dst, src net.Conn) {

		io.Copy(dst, src)

		dst.Close()
		wg.Done()
	}

	wg.Add(2)

	go cp(dst, src)
	go cp(src, dst)

	wg.Wait()
}

func server_mux(conn net.Conn) {
	defer conn.Close()

	sess, err := yamux.Server(conn, nil)
	if err != nil {
		log.Fatal(err)
	}
	defer sess.Close()

	for {
		src, err := sess.Accept()
		if err != nil {
			log.Println(err)
			break

		}

		if dst, err := net.Dial("tcp", *forward_addr); err == nil {
			go concat_conn(src, dst)

		} else {
			src.Close()
			log.Println(err)
		}
	}
}

func run_server() {

	log.Printf("server:\nlocal_addr:%s\nforward_addr:%s", *local_addr, *forward_addr)

	loop_accept(func(conn net.Conn) {
		go server_mux(conn)
	})
}

func run_client() {

	log.Printf("client:\nlocal_addr:%s\nremote_:%s", *local_addr, *remote_addr)

	var sess *yamux.Session

	new_session := func() error {
		log.Println("new session")
		conn, err := net.Dial("tcp", *remote_addr)
		if err != nil {
			log.Println(err)
			return err
		}

		sess, err = yamux.Client(conn, nil)
		if err != nil {
			log.Fatal(err)
		}

		return nil
	}

	loop_accept(func(src net.Conn) {

		if sess == nil {
			if err := new_session(); err != nil {
				src.Close()
				return
			}
		}

		dst, err := sess.Open()
		if err != nil {
			log.Println("stream open fail")
			sess.Close()
			sess = nil
			src.Close()
			return
		}

		go concat_conn(src, dst)

	})

}