基于 Yamux 的内网穿透简易实现

package main

import (
    "fmt"
    "log"
    "net/http"
)

func main() {

    http.HandleFunc("/hello", HelloHandler)
    fmt.Println("Server started at port 8881")
    log.Fatal(http.ListenAndServe(":8881", nil))
}

func HelloHandler(w http.ResponseWriter, r *http.Request) {
    fmt.Fprintf(w, "Hello, there\n")
}

type SessionManager interface {
	AddSession(key string, sess *yamux.Session)
	DialTarget(key string) (net.Conn, error)
}

func NewSessionManager() SessionManager {
	return &sessionManager{
		sessions: map[string][]*yamux.Session{},
	}
}

type sessionManager struct {
	sessions map[string][]*yamux.Session
	mutex    sync.Mutex
}

func (m *sessionManager) AddSession(key string, sess *yamux.Session) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	curr := m.sessions[key]
	if curr == nil {
		curr = []*yamux.Session{}
	}
	curr = append(curr, sess)
	m.sessions[key] = curr
}

func (m *sessionManager) DialTarget(key string) (net.Conn, error) {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	ss := m.sessions[key]
	for {
		if len(ss) == 0 {
			return nil, fmt.Errorf("no session found in '%s'", key)
		}

		idx := rand.Intn(len(ss))
		conn, err := ss[idx].Open()
		if err != nil {
			log.Printf("removing session #%d of '%s' due to dial error: %s", idx, key, err)
			ss[idx].Close()
			ss = append(ss[:idx], ss[idx+1:]...)
			m.sessions[key] = ss
			continue
		}
		log.Printf("find session with key '%s'", key)
		return conn, nil
	}
}

func SetupHandlers(r *gin.RouterGroup) {
	m := NewSessionManager()

	r.GET("/hubs/:id", func(c *gin.Context) {
		upgrader := websocket.Upgrader{
			ReadBufferSize:  1024,
			WriteBufferSize: 1024,
		}
		conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"msg": fmt.Sprintf("failed to upgrade to WebSocket: %s", err)})
			return
		}

		session, err := yamux.Server(conn.UnderlyingConn(), nil)
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"msg": fmt.Sprintf("failed to multiplex channel: %s", err)})
			return
		}

		m.AddSession(c.Param("id"), session)
	})

	r.Any("/proxy/:id/*proxyPath", func(c *gin.Context) {
		u, err := url.Parse("http://127.0.0.1")
		if err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"msg": err.Error()})
			return
		}

		rp := httputil.NewSingleHostReverseProxy(u)
		key := c.Param("id")
		rp.Transport = &http.Transport{
			DialContext: func(ctx context.Context, network string, addr string) (net.Conn, error) {
				conn, err := m.DialTarget(key)
				return conn, err
			},
		}
		rp.ErrorHandler = func(w http.ResponseWriter, r *http.Request, err error) {
			c.JSON(http.StatusBadRequest, gin.H{"msg": err.Error()})
		}

		hostInQuery, exist := c.GetQuery("x-proxy-host")
		if exist == true {
			c.Request.Header.Add("X-Proxy-Host", hostInQuery)
		}
		c.Request.Header.Add("X-Proxy-Path", c.Param("proxyPath"))
		rp.ServeHTTP(c.Writer, c.Request)
	})
}


func setupRouter() (*gin.Engine, error) {
	r := gin.Default()

	v1 := r.Group("/api/v1")
	SetupHandlers(v1)

	return r, nil
}

func Setup(connectURL string) error {
	dialer := websocket.DefaultDialer
	dialer.TLSClientConfig = &tls.Config{
		InsecureSkipVerify: true,
	}
	wsConn, _, err := dialer.Dial(connectURL, nil)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to dial hub %q: %s", connectURL, err))
	}

	sess, err := yamux.Client(wsConn.UnderlyingConn(), nil)
	if err != nil {
		return errors.New(fmt.Sprintf("failed to create multiplex channel: %s", err))
	}
	log.Println("connected to hub")

	director := func(req *http.Request) {
		log.Println(req)
		host := req.Header.Get("X-Proxy-Host")
		path := req.Header.Get("X-Proxy-Path")
		req.Header.Add("X-Forwarded-Host", req.Host)
		req.Header.Add("X-Origin-Host", host)
		req.URL.Scheme = "http"
		req.URL.Host = host
		req.URL.Path = path
		req.Host = ""
		log.Println(req)
	}

	proxy := &httputil.ReverseProxy{Director: director}
	server := &http.Server{
		Handler: proxy,
	}

	idleConnsClosed := make(chan struct{})
	go func() {
		sigint := make(chan os.Signal, 1)
		signal.Notify(sigint, os.Interrupt)
		signal.Notify(sigint, syscall.SIGTERM)
		<-sigint

		if err := server.Shutdown(context.Background()); err != nil {
			log.Printf("failed to shutdown server: %s", err)
		}
		close(idleConnsClosed)
	}()

	log.Println("starting proxy")
	if err := server.Serve(sess); err != nil {
		return errors.New(fmt.Sprintf("error running proxy: %s", err))
	}

	<-idleConnsClosed

	return nil
}