proxyport/app/forward/forward.go

package forward

import (
	"context"
	"fmt"
	"git.makemake.in/kzkzzzz/mycommon/mylog"
	"proxyport/app/db"
	"proxyport/app/model"
	"sync"
)

type Protocol int

func (p Protocol) String() string {
	switch p {
	case 0:
		return "TCP"
	case 1:
		return "UDP"
	}

	return "unknown"
}

const (
	ProtocolTCP Protocol = 0
	ProtocolUDP Protocol = 1
)

type IForward interface {
	Forward() error
	Stop()
}

var ListenerManager = &Manager{
	lock: &sync.Mutex{},
}

type Info struct {
	Id         int
	Name       string
	TargetAddr []string
	LocalAddr  string
	LocalIp    string
	LocalPort  int
	Protocol   Protocol
	Status     int
}

type activeForward struct {
	forward IForward
	info    *Info
}

type forwardKey struct {
	localAddr string
	protocol  Protocol
}

type Manager struct {
	activeForwardMap map[forwardKey]*activeForward
	lock             *sync.Mutex
}

func (m *Manager) initForwardList() []*Info {
	data := make([]*model.Forward, 0)
	err := db.DB().Table("forward").Select("*").Find(&data).Error
	if err != nil {
		panic(err)
	}

	list := make([]*Info, 0)

	for _, v := range data {
		if v.Status == 0 {
			continue
		}

		list = append(list, m.convertModel(v))
	}

	return list
}

func (m *Manager) convertModel(v *model.Forward) *Info {
	info := &Info{
		Id:         v.Id,
		Name:       v.Name,
		TargetAddr: v.TargetAddr,
		LocalIp:    v.LocalIp,
		LocalPort:  v.LocalPort,
		Protocol:   Protocol(v.Protocol),
		Status:     v.Status,
	}

	if v.LocalIp == "" {
		info.LocalAddr = fmt.Sprintf("0.0.0.0:%d", v.LocalPort)
	} else {
		info.LocalAddr = fmt.Sprintf("%s:%d", v.LocalIp, v.LocalPort)
	}

	return info
}

func (m *Manager) Add(v *model.Forward) error {
	m.lock.Lock()
	defer m.lock.Unlock()

	forwardInfo := m.convertModel(v)

	key := forwardKey{
		localAddr: forwardInfo.LocalAddr,
		protocol:  forwardInfo.Protocol,
	}

	_, ok := m.activeForwardMap[key]
	if ok {
		return fmt.Errorf("[%s] [%s] is exist", forwardInfo.Protocol, forwardInfo.LocalAddr)
	}

	tmp := &activeForward{
		info: forwardInfo,
	}
	if Protocol(v.Protocol) == ProtocolUDP {
		tmp.forward = NewUDP(forwardInfo)
	} else {
		tmp.forward = NewTCP(forwardInfo)
	}
	m.activeForwardMap[key] = tmp

	err := tmp.forward.Forward()
	if err != nil {
		delete(m.activeForwardMap, key)
		return err
	}

	return nil
}

func (m *Manager) Remove(v *model.Forward) {
	m.lock.Lock()
	defer m.lock.Unlock()

	forwardInfo := m.convertModel(v)

	key := forwardKey{
		localAddr: forwardInfo.LocalAddr,
		protocol:  forwardInfo.Protocol,
	}
	fr, ok := m.activeForwardMap[key]
	if !ok {
		return
	}

	fr.forward.Stop()
	delete(m.activeForwardMap, key)
}

func (m *Manager) Start(ctx context.Context) {
	m.activeForwardMap = make(map[forwardKey]*activeForward)

	list := m.initForwardList()
	for _, v := range list {
		key := forwardKey{
			localAddr: v.LocalAddr,
			protocol:  v.Protocol,
		}

		switch v.Protocol {
		case ProtocolTCP:
			m.activeForwardMap[key] = &activeForward{
				forward: NewTCP(v),
				info:    v,
			}

		case ProtocolUDP:
			m.activeForwardMap[key] = &activeForward{
				forward: NewUDP(v),
				info:    v,
			}
		}
	}

	for _, v := range m.activeForwardMap {
		err := v.forward.Forward()
		if err != nil {
			mylog.Error(err)
		}
	}

	select {
	case <-ctx.Done():
		m.Stop()
	}
}

func (m *Manager) Stop() {
	for _, v := range m.activeForwardMap {
		v.forward.Stop()
	}
}