proxy-lite/internal/app/app.go

package app

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

	"infinite-noodle/internal/noodle"
	"infinite-noodle/internal/web"

	"inet.af/tcpproxy"
)

type Config struct {
	DataPath string
	Host     string
	Port     int
	RunTest  bool
}

func Run(cfg Config) error {
	listenAddr := fmt.Sprintf("%s:%d", cfg.Host, cfg.Port)
	noodleChannel := make(chan noodle.Noodle)

	db := noodle.NewDatabase(cfg.DataPath)
	defer db.Close()

	if cfg.RunTest {
		runTestSequence(db)
	}

	go systemCheck(db, noodleChannel)
	go expirationCheck(db, noodleChannel)
	go tcpProxify(noodleChannel)

	http.Handle("/static/", http.StripPrefix("/static/", http.FileServer(http.FS(web.StaticFiles()))))
	http.HandleFunc("/", web.HandleMain(db, &noodleChannel))
	http.HandleFunc("/add", web.HandleAdd(db, &noodleChannel))
	http.HandleFunc("/toggle", web.HandleToggle(db, &noodleChannel))
	http.HandleFunc("/delete", web.HandleDelete(db, &noodleChannel))

	log.Printf("Server starting on %s", listenAddr)
	return http.ListenAndServe(listenAddr, nil)
}

func systemCheck(db *noodle.Database, noodleChannel chan noodle.Noodle) {
	for {
		noodles := db.GetAll()
		for _, item := range noodles {
			noodleChannel <- item
		}
		time.Sleep(60 * time.Second)
	}
}

func startProxy(proxy *tcpproxy.Proxy) {
	log.Print(proxy.Run())
}

func tcpProxify(noodleChannel chan noodle.Noodle) {
	noodleMap := make(map[string]*tcpproxy.Proxy)
	for {
		item := <-noodleChannel
		_, running := noodleMap[item.Id]
		if item.IsUp && !running {
			var proxy tcpproxy.Proxy
			src := fmt.Sprintf("0.0.0.0:%d", item.ListenPort)
			dst := fmt.Sprintf("%s:%d", item.DestHost, item.DestPort)
			log.Printf("Starting a noodle from %s to %s with source=%s", src, dst, item.Src)
			proxy.AddRoute(src, sourceRestrictedTarget{
				allowedIP: item.Src,
				target:    tcpproxy.To(dst),
			})
			noodleMap[item.Id] = &proxy
			go startProxy(&proxy)
			continue
		}

		if !item.IsUp && running {
			log.Printf("Closing noodle=%v", item)
			if err := noodleMap[item.Id].Close(); err != nil {
				log.Print(err)
			}
			delete(noodleMap, item.Id)
		}
	}
}

type sourceRestrictedTarget struct {
	allowedIP string
	target    tcpproxy.Target
}

func (t sourceRestrictedTarget) HandleConn(conn net.Conn) {
	if t.allowedIP == "" || t.allowedIP == "All" {
		t.target.HandleConn(conn)
		return
	}

	host, _, err := net.SplitHostPort(conn.RemoteAddr().String())
	if err != nil {
		log.Printf("Rejected noodle connection with invalid remote address %q", conn.RemoteAddr().String())
		conn.Close()
		return
	}
	if host != t.allowedIP {
		log.Printf("Rejected noodle connection from %s; allowed source is %s", host, t.allowedIP)
		conn.Close()
		return
	}

	t.target.HandleConn(conn)
}

func expirationCheck(db *noodle.Database, noodleChannel chan noodle.Noodle) {
	ticker := time.NewTicker(1 * time.Second)
	defer ticker.Stop()

	for range ticker.C {
		noodles := db.GetAll()
		for _, item := range noodles {
			if !item.IsUp {
				continue
			}
			if item.Expiration <= 0 {
				if item.IsUp {
					item.IsUp = false
					noodleChannel <- item
				}
				if err := db.Delete(item.Id); err != nil {
					log.Print(err)
				}
				continue
			}

			item.Expiration -= time.Second
			if item.Expiration <= 0 {
				item.Expiration = 0
				if err := db.Update(item); err != nil {
					log.Print(err)
					continue
				}
				if item.IsUp {
					item.IsUp = false
					noodleChannel <- item
				}
				if err := db.Delete(item.Id); err != nil {
					log.Print(err)
				}
				continue
			}

			if err := db.Update(item); err != nil {
				log.Print(err)
			}
		}
	}
}

func runTestSequence(db *noodle.Database) {
	for i := 0; i < 21; i++ {
		item := noodle.Noodle{
			Id:         db.MakeID(),
			Name:       "Name_Test",
			Proto:      "Proto_Test",
			Src:        "All",
			ListenPort: 1080 + i,
			DestPort:   22,
			DestHost:   "localhost",
			Expiration: time.Duration(time.Now().Second()) * time.Second,
			IsUp:       true,
		}
		log.Printf("Test noodle=%v", item)
		db.Add(item)
		log.Printf("Test id=%s exists=%t", item.Id, db.Handle.Has(item.Id))
		log.Printf("Test GetAllGeneric=%v", db.GetAllGeneric())
	}
}