init

at-server/structs.go

@@ -0,0 +1,231 @@
+package main
+
+import (
+	utils "anytunnel/at-common"
+	"crypto/tls"
+	"fmt"
+	"net"
+	"strconv"
+	"time"
+)
+
+//###############ServerUserConn##################
+
+type ServerUserConn struct {
+	TunnelID     uint64
+	ConnectionID uint64
+	Connection   *tls.Conn
+}
+
+type ServerUserConnPool struct {
+	pool utils.ConcurrentMap
+}
+
+func NewServerUserConnPool() ServerUserConnPool {
+	return ServerUserConnPool{
+		pool: utils.NewConcurrentMap(),
+	}
+}
+func (m *ServerUserConnPool) IntToString(id uint64) string {
+	return strconv.FormatUint(id, 10)
+}
+func (t *ServerUserConnPool) Put(tunnelID, connectionID uint64, conn *tls.Conn) {
+	tid := t.IntToString(tunnelID)
+	cid := t.IntToString(connectionID)
+	if _, e := t.pool.Get(tid); !e {
+		mp := utils.NewConcurrentMap()
+		t.pool.Set(tid, &mp)
+	}
+	m, _ := t.pool.Get(tid)
+	m.(*utils.ConcurrentMap).Set(cid, ServerUserConn{
+		Connection:   conn,
+		TunnelID:     tunnelID,
+		ConnectionID: tunnelID,
+	})
+}
+func (t *ServerUserConnPool) Get(tunnelID, connectionID uint64) (conn ServerUserConn, err error) {
+	tid := t.IntToString(tunnelID)
+	cid := t.IntToString(connectionID)
+	_, e := t.pool.Get(tid)
+	if !e {
+		err = fmt.Errorf("tunnel %d not exists", tunnelID)
+		return
+	}
+	m, _ := t.pool.Get(tid)
+
+	v1, e1 := m.(*utils.ConcurrentMap).Get(cid)
+	if !e1 {
+		err = fmt.Errorf("server user connection %d not exists", connectionID)
+		return
+	}
+	return v1.(ServerUserConn), err
+}
+func (t *ServerUserConnPool) Exists(tunnelID uint64, connectionID uint64) bool {
+	tid := t.IntToString(tunnelID)
+	cid := t.IntToString(connectionID)
+	_, e := t.pool.Get(tid)
+	if e {
+		m, _ := t.pool.Get(tid)
+		_, e = m.(*utils.ConcurrentMap).Get(cid)
+	}
+	return e
+}
+func (t *ServerUserConnPool) DeleteAll(tunnelID uint64) {
+	tid := t.IntToString(tunnelID)
+	_, e := t.pool.Get(tid)
+	if !e {
+		return
+	}
+	m, _ := t.pool.Get(tid)
+	for _, v := range m.(*utils.ConcurrentMap).Keys() {
+		el, _ := m.(*utils.ConcurrentMap).Get(v)
+		(*(el.(ServerUserConn)).Connection).Close()
+	}
+	t.pool.Remove(tid)
+}
+func (t *ServerUserConnPool) Delete(tunnelID uint64, connectionID uint64) {
+	tid := t.IntToString(tunnelID)
+	cid := t.IntToString(connectionID)
+	v1, e := t.pool.Get(tid)
+	if !e {
+		return
+	}
+	c, e := v1.(*utils.ConcurrentMap).Get(cid)
+	if !e {
+		return
+	}
+	(*(c.(ServerUserConn)).Connection).Close()
+	v1.(*utils.ConcurrentMap).Remove(cid)
+	if v1.(*utils.ConcurrentMap).IsEmpty() {
+		t.pool.Remove(tid)
+	}
+}
+
+//###############ListenerMap##################
+type ListenerMap struct {
+	data utils.ConcurrentMap
+}
+type ListenerTimeout struct {
+	Listener *net.Listener
+	LastTime int64
+}
+type UDPListenerTimeout struct {
+	Listener *net.UDPConn
+	LastTime int64
+}
+
+func NewListenerMap() ListenerMap {
+	cm := ListenerMap{
+		data: utils.NewConcurrentMap(),
+	}
+	return cm
+}
+func (m *ListenerMap) IntToString(id uint64) string {
+	return strconv.FormatUint(id, 10)
+}
+func (m *ListenerMap) Get(id uint64) (conn *net.Listener, exists bool) {
+	mid := m.IntToString(id)
+	_conn, exists := m.data.Get(mid)
+	if exists {
+		conn = (_conn).(ListenerTimeout).Listener
+	}
+	return
+}
+func (m *ListenerMap) GetUDP(id uint64) (conn *net.UDPConn, exists bool) {
+	mid := m.IntToString(id)
+	_conn, exists := m.data.Get(mid)
+	if exists {
+		conn = (_conn).(UDPListenerTimeout).Listener
+	}
+	return
+}
+func (m *ListenerMap) Put(id uint64, Listener interface{}) {
+	mid := m.IntToString(id)
+	if v, ok := Listener.(*net.Listener); ok {
+		m.data.Set(mid, ListenerTimeout{
+			Listener: v,
+			LastTime: time.Now().Unix(),
+		})
+	} else {
+		m.data.Set(mid, UDPListenerTimeout{
+			Listener: Listener.(*net.UDPConn),
+			LastTime: time.Now().Unix(),
+		})
+	}
+
+	return
+}
+func (m *ListenerMap) Delete(id uint64) {
+	mid := m.IntToString(id)
+	if _conn, exists := m.data.Get(mid); exists {
+		if _, ok := _conn.(ListenerTimeout); ok {
+			(*(_conn.(ListenerTimeout)).Listener).Close()
+		} else {
+			(*(_conn.(UDPListenerTimeout)).Listener).Close()
+		}
+
+	}
+	m.data.Remove(mid)
+}
+
+//#################ip conn counter################
+type IPConnCounter struct {
+	data         utils.ConcurrentMap
+	blockSeconds int
+	maxConnCount int
+}
+type IPConnCounterItem struct {
+	LastTime int64
+	Count    int
+}
+
+func NewIPConnCounter(maxConnCount, blockSeconds int) (ic IPConnCounter) {
+	ic = IPConnCounter{
+		data:         utils.NewConcurrentMap(),
+		blockSeconds: blockSeconds,
+		maxConnCount: maxConnCount - 1,
+	}
+	ic.gc()
+	return
+}
+func (ic *IPConnCounter) Check(ip string) bool {
+	_item, ok := ic.data.Get(ip)
+	if ok {
+		item := _item.(IPConnCounterItem)
+		if item.Count > ic.maxConnCount {
+			if time.Now().Unix()-item.LastTime > int64(ic.blockSeconds) {
+				ic.data.Remove(ip)
+				return true
+			}
+			return false
+		}
+	}
+	ic.data.Upsert(ip, nil, func(exist bool, valueInMap interface{}, newValue interface{}) (res interface{}) {
+		var _valueInMap IPConnCounterItem
+		if exist {
+			_valueInMap = valueInMap.(IPConnCounterItem)
+			_valueInMap.Count = _valueInMap.Count + 1
+			_valueInMap.LastTime = time.Now().Unix()
+		} else {
+			_valueInMap = IPConnCounterItem{
+				Count:    1,
+				LastTime: time.Now().Unix(),
+			}
+		}
+		return _valueInMap
+	})
+	return true
+}
+func (ic *IPConnCounter) gc() {
+	go func() {
+		for {
+			for k, v := range ic.data.Items() {
+				item := v.(IPConnCounterItem)
+				if time.Now().Unix()-item.LastTime > int64(ic.blockSeconds) {
+					ic.data.Remove(k)
+				}
+			}
+			time.Sleep(time.Second * 300)
+		}
+	}()
+}