EdgeNode/internal/nodes/client_conn_base.go

// Copyright 2021 GoEdge goedge.cdn@gmail.com. All rights reserved.

package nodes

import (
	"crypto/tls"
	"net"
	"sync/atomic"
	"time"

	"github.com/TeaOSLab/EdgeNode/internal/firewalls"
	"github.com/TeaOSLab/EdgeNode/internal/iplibrary"
)

type BaseClientConn struct {
	rawConn net.Conn

	isBound    bool
	userId     int64
	userPlanId int64
	serverId   int64
	remoteAddr string
	hasLimit   bool

	isPersistent bool // 是否为持久化连接
	fingerprint  []byte

	isClosed bool

	rawIP string

	totalSentBytes int64
}

func (this *BaseClientConn) IsClosed() bool {
	return this.isClosed
}

// IsBound 是否已绑定服务
func (this *BaseClientConn) IsBound() bool {
	return this.isBound
}

// Bind 绑定服务
func (this *BaseClientConn) Bind(serverId int64, remoteAddr string, maxConnsPerServer int, maxConnsPerIP int) bool {
	if this.isBound {
		return true
	}
	this.isBound = true
	this.serverId = serverId
	this.remoteAddr = remoteAddr
	this.hasLimit = true

	// 检查是否可以连接
	return sharedClientConnLimiter.Add(this.rawConn.RemoteAddr().String(), serverId, remoteAddr, maxConnsPerServer, maxConnsPerIP)
}

// SetServerId 设置服务ID
func (this *BaseClientConn) SetServerId(serverId int64) (goNext bool) {
	goNext = true

	// 检查服务相关IP黑名单
	var rawIP = this.RawIP()
	if serverId > 0 && len(rawIP) > 0 {
		// 是否在白名单中
		ok, _, expiresAt := iplibrary.AllowIP(rawIP, serverId)
		if !ok {
			_ = this.rawConn.Close()
			firewalls.DropTemporaryTo(rawIP, expiresAt)
			return false
		}
	}

	this.serverId = serverId

	// 设置包装前连接
	switch conn := this.rawConn.(type) {
	case *tls.Conn:
		nativeConn, ok := conn.NetConn().(ClientConnInterface)
		if ok {
			nativeConn.SetServerId(serverId)
		}
	case *ClientConn:
		conn.SetServerId(serverId)
	}

	return true
}

// ServerId 读取当前连接绑定的服务ID
func (this *BaseClientConn) ServerId() int64 {
	return this.serverId
}

// SetUserId 设置所属服务的用户ID
func (this *BaseClientConn) SetUserId(userId int64) {
	this.userId = userId

	// 设置包装前连接
	switch conn := this.rawConn.(type) {
	case *tls.Conn:
		nativeConn, ok := conn.NetConn().(ClientConnInterface)
		if ok {
			nativeConn.SetUserId(userId)
		}
	case *ClientConn:
		conn.SetUserId(userId)
	}
}

func (this *BaseClientConn) SetUserPlanId(userPlanId int64) {
	this.userPlanId = userPlanId

	// 设置包装前连接
	switch conn := this.rawConn.(type) {
	case *tls.Conn:
		nativeConn, ok := conn.NetConn().(ClientConnInterface)
		if ok {
			nativeConn.SetUserPlanId(userPlanId)
		}
	case *ClientConn:
		conn.SetUserPlanId(userPlanId)
	}
}

// UserId 获取当前连接所属服务的用户ID
func (this *BaseClientConn) UserId() int64 {
	return this.userId
}

// UserPlanId 用户套餐ID
func (this *BaseClientConn) UserPlanId() int64 {
	return this.userPlanId
}

// RawIP 原本IP
func (this *BaseClientConn) RawIP() string {
	if len(this.rawIP) > 0 {
		return this.rawIP
	}

	ip, _, _ := net.SplitHostPort(this.rawConn.RemoteAddr().String())
	this.rawIP = ip
	return ip
}

// TCPConn 转换为TCPConn
func (this *BaseClientConn) TCPConn() (tcpConn *net.TCPConn, ok bool) {
	// 设置包装前连接
	switch conn := this.rawConn.(type) {
	case *tls.Conn:
		var internalConn = conn.NetConn()
		clientConn, isClientConn := internalConn.(*ClientConn)
		if isClientConn {
			return clientConn.TCPConn()
		}
		tcpConn, ok = internalConn.(*net.TCPConn)
	default:
		tcpConn, ok = this.rawConn.(*net.TCPConn)
	}
	return
}

// SetLinger 设置Linger
func (this *BaseClientConn) SetLinger(seconds int) error {
	tcpConn, ok := this.TCPConn()
	if ok {
		return tcpConn.SetLinger(seconds)
	}
	return nil
}

func (this *BaseClientConn) SetIsPersistent(isPersistent bool) {
	this.isPersistent = isPersistent

	_ = this.rawConn.SetDeadline(time.Time{})
}

// SetFingerprint 设置指纹信息
func (this *BaseClientConn) SetFingerprint(fingerprint []byte) {
	this.fingerprint = fingerprint
}

// Fingerprint 读取指纹信息
func (this *BaseClientConn) Fingerprint() []byte {
	return this.fingerprint
}

// LastRequestBytes 读取上一次请求发送的字节数
func (this *BaseClientConn) LastRequestBytes() int64 {
	var result = atomic.LoadInt64(&this.totalSentBytes)
	atomic.StoreInt64(&this.totalSentBytes, 0)
	return result
}