added unlinked controller; change create forwarded to createOrUpdate; added global proxy mutex
This commit is contained in:
@@ -0,0 +1,25 @@
|
|||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
kubeclient "k8s.io/client-go/kubernetes"
|
||||||
|
kubeclicmd "k8s.io/client-go/tools/clientcmd"
|
||||||
|
)
|
||||||
|
|
||||||
|
func makeClientset(kubeconf []byte) (kubeclient.Interface, error) {
|
||||||
|
var res kubeclient.Interface
|
||||||
|
var err error
|
||||||
|
clientConfig, err := kubeclicmd.NewClientConfigFromBytes(kubeconf)
|
||||||
|
if err != nil {
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
restConfig, err := clientConfig.ClientConfig()
|
||||||
|
if err != nil {
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
kubeClient, err := kubeclient.NewForConfig(restConfig)
|
||||||
|
if err != nil {
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
|
res = kubeClient
|
||||||
|
return res, err
|
||||||
|
}
|
||||||
@@ -0,0 +1,226 @@
|
|||||||
|
package control
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"encoding/json"
|
||||||
|
"time"
|
||||||
|
|
||||||
|
kubecore "k8s.io/api/core/v1"
|
||||||
|
k8inform "k8s.io/client-go/informers"
|
||||||
|
k8client "k8s.io/client-go/kubernetes"
|
||||||
|
k8cache "k8s.io/client-go/tools/cache"
|
||||||
|
|
||||||
|
kubemeta "k8s.io/apimachinery/pkg/apis/meta/v1"
|
||||||
|
kubetypes "k8s.io/apimachinery/pkg/types"
|
||||||
|
kubepatch "k8s.io/apimachinery/pkg/util/strategicpatch"
|
||||||
|
|
||||||
|
"helmet/app/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Controller struct {
|
||||||
|
lbaddr string
|
||||||
|
clientset k8client.Interface
|
||||||
|
informer k8cache.SharedIndexInformer
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
log *logger.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewController(clientset k8client.Interface, lbaddr string) *Controller {
|
||||||
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
|
cont := &Controller{
|
||||||
|
clientset: clientset,
|
||||||
|
ctx: ctx,
|
||||||
|
cancel: cancel,
|
||||||
|
lbaddr: lbaddr,
|
||||||
|
}
|
||||||
|
cont.log = logger.NewLogger("controller")
|
||||||
|
return cont
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) Run() {
|
||||||
|
cont.log.Debugf("Start controller")
|
||||||
|
factory := k8inform.NewSharedInformerFactory(cont.clientset, time.Minute*10)
|
||||||
|
defer factory.Shutdown()
|
||||||
|
|
||||||
|
serviceInformer := factory.Core().V1().Services().Informer()
|
||||||
|
handler := k8cache.ResourceEventHandlerFuncs{
|
||||||
|
AddFunc: cont.addService,
|
||||||
|
UpdateFunc: cont.updateService,
|
||||||
|
DeleteFunc: cont.deleteService,
|
||||||
|
}
|
||||||
|
serviceInformer.AddEventHandler(handler)
|
||||||
|
|
||||||
|
ctx, cancel := context.WithCancel(cont.ctx)
|
||||||
|
defer cancel()
|
||||||
|
factory.Start(ctx.Done())
|
||||||
|
synced := factory.WaitForCacheSync(ctx.Done())
|
||||||
|
for _, sync := range synced {
|
||||||
|
if !sync {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
cont.informer = serviceInformer
|
||||||
|
<-ctx.Done()
|
||||||
|
cont.log.Debugf("Stop controller")
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) Stop() {
|
||||||
|
cont.cancel()
|
||||||
|
}
|
||||||
|
|
||||||
|
// https://pkg.go.dev/k8s.io/api/core/v1#Service
|
||||||
|
|
||||||
|
func (cont *Controller) addService(obj any) {
|
||||||
|
service := obj.(*kubecore.Service)
|
||||||
|
service = service.DeepCopy()
|
||||||
|
cont.log.Debugf("Service %s/%s created", service.Namespace, service.Name)
|
||||||
|
if service.Spec.Type == kubecore.ServiceTypeLoadBalancer {
|
||||||
|
_ = makeForwarding(service)
|
||||||
|
err := cont.patchService(service)
|
||||||
|
if err != nil {
|
||||||
|
cont.log.Debugf("Error patch service %s/%s: %v", service.Namespace, service.Name, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) updateService(oldObj, newObj any) {
|
||||||
|
newService := newObj.(*kubecore.Service)
|
||||||
|
newService = newService.DeepCopy()
|
||||||
|
oldService := oldObj.(*kubecore.Service)
|
||||||
|
oldService = oldService.DeepCopy()
|
||||||
|
cont.log.Debugf("Service %s/%s updated", newService.Namespace, newService.Name)
|
||||||
|
|
||||||
|
if newService.Spec.Type == kubecore.ServiceTypeLoadBalancer {
|
||||||
|
if cont.serviceDifferent(oldService, newService) {
|
||||||
|
_ = makeForwarding(newService)
|
||||||
|
cont.log.Debugf("Updated service %s/%s have new adresses", newService.Namespace, newService.Name)
|
||||||
|
err := cont.patchService(newService)
|
||||||
|
if err != nil {
|
||||||
|
cont.log.Debugf("Error patch service %s/%s: %v", newService.Namespace, newService.Name, err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) deleteService(obj any) {
|
||||||
|
service := obj.(*kubecore.Service)
|
||||||
|
service = service.DeepCopy()
|
||||||
|
cont.log.Debugf("Service %s/%s deleted", service.Namespace, service.Name)
|
||||||
|
}
|
||||||
|
|
||||||
|
type Forwarding struct {
|
||||||
|
Proto string
|
||||||
|
Port uint32
|
||||||
|
Destinations []string
|
||||||
|
}
|
||||||
|
|
||||||
|
func makeForwarding(service *kubecore.Service) []*Forwarding {
|
||||||
|
forwardings := make([]*Forwarding, 0)
|
||||||
|
for _, port := range service.Spec.Ports {
|
||||||
|
forwarding := &Forwarding{
|
||||||
|
Proto: string(port.Protocol),
|
||||||
|
Port: uint32(port.Port),
|
||||||
|
Destinations: make([]string, 0),
|
||||||
|
}
|
||||||
|
for _, ipaddr := range service.Spec.ClusterIPs {
|
||||||
|
forwarding.Destinations = append(forwarding.Destinations, ipaddr)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return forwardings
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) patchService(svc *kubecore.Service) error {
|
||||||
|
var err error
|
||||||
|
oldData, err := json.Marshal(svc)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
ingressMode := kubecore.LoadBalancerIPModeProxy
|
||||||
|
ingresses := make([]kubecore.LoadBalancerIngress, 0)
|
||||||
|
ingresses = append(ingresses, kubecore.LoadBalancerIngress{
|
||||||
|
IP: cont.lbaddr,
|
||||||
|
IPMode: &ingressMode,
|
||||||
|
})
|
||||||
|
status := kubecore.LoadBalancerStatus{
|
||||||
|
Ingress: ingresses,
|
||||||
|
}
|
||||||
|
svc.Status.LoadBalancer = status
|
||||||
|
|
||||||
|
newData, err := json.Marshal(svc)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
patchBytes, err := kubepatch.CreateTwoWayMergePatch(oldData, newData, kubecore.Service{})
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
cont.log.Debugf("patch: %s\n", string(patchBytes))
|
||||||
|
|
||||||
|
ctx, _ := context.WithTimeout(cont.ctx, 5*time.Second)
|
||||||
|
patchServiceOpts := kubemeta.PatchOptions{}
|
||||||
|
strategy := kubetypes.StrategicMergePatchType
|
||||||
|
_, err = cont.clientset.CoreV1().Services(svc.Namespace).
|
||||||
|
Patch(ctx, svc.Name, strategy, patchBytes, patchServiceOpts, "status")
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (cont *Controller) serviceDifferent(oldService, newService *kubecore.Service) bool {
|
||||||
|
type Address struct {
|
||||||
|
Port int32
|
||||||
|
Proto kubecore.Protocol
|
||||||
|
Address string
|
||||||
|
}
|
||||||
|
oldAddreses := make([]Address, 0)
|
||||||
|
for _, port := range oldService.Spec.Ports {
|
||||||
|
for _, ipaddr := range oldService.Spec.ClusterIPs {
|
||||||
|
address := Address{
|
||||||
|
Port: port.Port,
|
||||||
|
Proto: port.Protocol,
|
||||||
|
Address: ipaddr,
|
||||||
|
}
|
||||||
|
oldAddreses = append(oldAddreses, address)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
newAddreses := make([]Address, 0)
|
||||||
|
for _, port := range newService.Spec.Ports {
|
||||||
|
for _, ipaddr := range newService.Spec.ClusterIPs {
|
||||||
|
address := Address{
|
||||||
|
Port: port.Port,
|
||||||
|
Proto: port.Protocol,
|
||||||
|
Address: ipaddr,
|
||||||
|
}
|
||||||
|
newAddreses = append(newAddreses, address)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if len(newAddreses) != len(oldAddreses) {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
for _, newAddress := range newAddreses {
|
||||||
|
newFound := false
|
||||||
|
for _, oldAddress := range oldAddreses {
|
||||||
|
if oldAddress == newAddress {
|
||||||
|
newFound = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !newFound {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
for _, oldAddress := range oldAddreses {
|
||||||
|
oldFound := false
|
||||||
|
for _, newAddress := range newAddreses {
|
||||||
|
if newAddress == oldAddress {
|
||||||
|
oldFound = true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if !oldFound {
|
||||||
|
return true
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return false
|
||||||
|
}
|
||||||
@@ -37,7 +37,7 @@ func NewLogger(subj string) *Logger {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func xxxNewLogger() *Logger {
|
func NewUnamedLogger() *Logger {
|
||||||
return &Logger{
|
return &Logger{
|
||||||
writer: output,
|
writer: output,
|
||||||
mtx: &mtx,
|
mtx: &mtx,
|
||||||
|
|||||||
@@ -12,7 +12,8 @@ func (oper *Operator) ListForwarders(ctx context.Context, params *mlbctl.ListFor
|
|||||||
res := &mlbctl.ListForwardersResult{
|
res := &mlbctl.ListForwardersResult{
|
||||||
Forwarders: make([]*mlbctl.Forwarder, 0),
|
Forwarders: make([]*mlbctl.Forwarder, 0),
|
||||||
}
|
}
|
||||||
for _, forw := range oper.proxy.Forwarders {
|
forws := oper.proxy.ListForwarders(ctx)
|
||||||
|
for _, forw := range forws {
|
||||||
oForw := &mlbctl.Forwarder{
|
oForw := &mlbctl.Forwarder{
|
||||||
Type: forw.Type,
|
Type: forw.Type,
|
||||||
Lport: forw.Lport,
|
Lport: forw.Lport,
|
||||||
@@ -33,8 +34,8 @@ func (oper *Operator) ListForwarders(ctx context.Context, params *mlbctl.ListFor
|
|||||||
func (oper *Operator) CreateForwarder(ctx context.Context, params *mlbctl.CreateForwarderParams) (*mlbctl.CreateForwarderResult, error) {
|
func (oper *Operator) CreateForwarder(ctx context.Context, params *mlbctl.CreateForwarderParams) (*mlbctl.CreateForwarderResult, error) {
|
||||||
var err error
|
var err error
|
||||||
res := &mlbctl.CreateForwarderResult{}
|
res := &mlbctl.CreateForwarderResult{}
|
||||||
err = oper.proxy.AddForwarder(ctx, params.Type, params.Lport, params.Dport, params.Destinations...)
|
err = oper.proxy.CreateOrUpdateForwarder(ctx, params.Type, params.Lport, params.Dport, params.Destinations...)
|
||||||
if err != err {
|
if err != nil {
|
||||||
return res, err
|
return res, err
|
||||||
}
|
}
|
||||||
return res, err
|
return res, err
|
||||||
@@ -43,8 +44,8 @@ func (oper *Operator) CreateForwarder(ctx context.Context, params *mlbctl.Create
|
|||||||
func (oper *Operator) DeleteForwarder(ctx context.Context, params *mlbctl.DeleteForwarderParams) (*mlbctl.DeleteForwarderResult, error) {
|
func (oper *Operator) DeleteForwarder(ctx context.Context, params *mlbctl.DeleteForwarderParams) (*mlbctl.DeleteForwarderResult, error) {
|
||||||
var err error
|
var err error
|
||||||
res := &mlbctl.DeleteForwarderResult{}
|
res := &mlbctl.DeleteForwarderResult{}
|
||||||
err = oper.proxy.DeleteForwarder(ctx, params.Lport)
|
err = oper.proxy.DeleteForwarder(ctx, params.Type, params.Lport)
|
||||||
if err != err {
|
if err != nil {
|
||||||
return res, err
|
return res, err
|
||||||
}
|
}
|
||||||
return res, err
|
return res, err
|
||||||
|
|||||||
+143
-87
@@ -2,11 +2,12 @@ package rproxy
|
|||||||
|
|
||||||
import (
|
import (
|
||||||
"context"
|
"context"
|
||||||
"io"
|
"fmt"
|
||||||
"math/rand"
|
"math/rand"
|
||||||
"net"
|
"net"
|
||||||
"strconv"
|
"strconv"
|
||||||
"sync"
|
"sync"
|
||||||
|
"time"
|
||||||
|
|
||||||
"helmet/app/logger"
|
"helmet/app/logger"
|
||||||
)
|
)
|
||||||
@@ -14,20 +15,27 @@ import (
|
|||||||
const (
|
const (
|
||||||
TCP = "tcp"
|
TCP = "tcp"
|
||||||
UDP = "udp"
|
UDP = "udp"
|
||||||
|
|
||||||
|
ForwStarted = "started"
|
||||||
|
ForwStopped = "stopped"
|
||||||
)
|
)
|
||||||
|
|
||||||
type Forwarder struct {
|
type Forwarder struct {
|
||||||
Type string `json:"type" yaml:"type"`
|
State string `json:"state" yaml:"state"`
|
||||||
listen net.Listener `json:"-" yaml:"-"`
|
Type string `json:"type" yaml:"type"`
|
||||||
ctx context.Context `json:"-" yaml:"-"`
|
Lport uint32 `json:"lport" yaml:"lport"`
|
||||||
cancel context.CancelFunc `json:"-" yaml:"-"`
|
Dport uint32 `json:"dport" yaml:"dport"`
|
||||||
Lport uint32 `json:"lport" yaml:"lport"`
|
Dests []*Destination `json:"dests" yaml:"dests"`
|
||||||
Dport uint32 `json:"dport" yaml:"dport"`
|
|
||||||
Dests []*Destination `json:"dests" yaml:"dests"`
|
listenTCP *net.TCPListener `json:"-" yaml:"-"`
|
||||||
log *logger.Logger
|
listenUDP *net.UDPConn `json:"-" yaml:"-"`
|
||||||
|
ctx context.Context `json:"-" yaml:"-"`
|
||||||
|
cancel context.CancelFunc `json:"-" yaml:"-"`
|
||||||
|
log *logger.Logger
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ...string) (*Forwarder, error) {
|
func NewForwarder(ctx context.Context, proto string, lport, dport uint32, addrs ...string) (*Forwarder, error) {
|
||||||
|
var err error
|
||||||
ctx, cancel := context.WithCancel(context.Background())
|
ctx, cancel := context.WithCancel(context.Background())
|
||||||
forw := &Forwarder{
|
forw := &Forwarder{
|
||||||
Dests: make([]*Destination, 0),
|
Dests: make([]*Destination, 0),
|
||||||
@@ -35,33 +43,67 @@ func NewForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ..
|
|||||||
Dport: dport,
|
Dport: dport,
|
||||||
ctx: ctx,
|
ctx: ctx,
|
||||||
cancel: cancel,
|
cancel: cancel,
|
||||||
Type: typ,
|
Type: proto,
|
||||||
}
|
}
|
||||||
id := strconv.FormatUint(uint64(lport), 10)
|
id := forw.Type + strconv.FormatUint(uint64(lport), 10)
|
||||||
forw.log = logger.NewLogger("forwarder:" + id)
|
forw.log = logger.NewLogger("forwarder:" + id)
|
||||||
for _, addr := range addrs {
|
for _, addr := range addrs {
|
||||||
dest := NewDestination(addr)
|
dest := NewDestination(addr)
|
||||||
forw.Dests = append(forw.Dests, dest)
|
forw.Dests = append(forw.Dests, dest)
|
||||||
}
|
}
|
||||||
|
switch proto {
|
||||||
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
case TCP:
|
||||||
laddr, err := net.ResolveTCPAddr("tcp", portinfo)
|
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
||||||
if err != nil {
|
laddr, err := net.ResolveTCPAddr("tcp", portinfo)
|
||||||
|
if err != nil {
|
||||||
|
return forw, err
|
||||||
|
}
|
||||||
|
listen, err := net.ListenTCP("tcp", laddr)
|
||||||
|
if err != nil {
|
||||||
|
return forw, err
|
||||||
|
}
|
||||||
|
forw.listenTCP = listen
|
||||||
|
case UDP:
|
||||||
|
portinfo := ":" + strconv.FormatUint(uint64(forw.Lport), 10)
|
||||||
|
laddr, err := net.ResolveUDPAddr("udp", portinfo)
|
||||||
|
if err != nil {
|
||||||
|
return forw, err
|
||||||
|
}
|
||||||
|
listen, err := net.ListenUDP("udp", laddr)
|
||||||
|
if err != nil {
|
||||||
|
return forw, err
|
||||||
|
}
|
||||||
|
forw.listenUDP = listen
|
||||||
|
default:
|
||||||
|
err = fmt.Errorf("Unknown net type: %s", proto)
|
||||||
return forw, err
|
return forw, err
|
||||||
}
|
}
|
||||||
listen, err := net.ListenTCP("tcp", laddr)
|
|
||||||
if err != nil {
|
|
||||||
return forw, err
|
|
||||||
}
|
|
||||||
forw.listen = listen
|
|
||||||
return forw, err
|
return forw, err
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (forw *Forwarder) Listen(wg *sync.WaitGroup) {
|
||||||
|
switch forw.Type {
|
||||||
|
case TCP:
|
||||||
|
if forw.listenTCP != nil {
|
||||||
|
forw.ListenTCP(wg)
|
||||||
|
}
|
||||||
|
case UDP:
|
||||||
|
if forw.listenUDP != nil {
|
||||||
|
forw.ListenUDP(wg)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
|
func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
|
||||||
forw.log.Debugf("Start listening on %d", forw.Lport)
|
forw.log.Debugf("Start listening on %s:%d", forw.Type, forw.Lport)
|
||||||
|
forw.State = ForwStarted
|
||||||
defer wg.Done()
|
defer wg.Done()
|
||||||
|
stater := func() {
|
||||||
|
forw.State = ForwStopped
|
||||||
|
}
|
||||||
|
defer stater()
|
||||||
for {
|
for {
|
||||||
conn, err := forw.listen.Accept()
|
conn, err := forw.listenTCP.Accept()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
forw.log.Errorf("Listen err: %v", err)
|
forw.log.Errorf("Listen err: %v", err)
|
||||||
return
|
return
|
||||||
@@ -70,8 +112,84 @@ func (forw *Forwarder) ListenTCP(wg *sync.WaitGroup) {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func (forw *Forwarder) ListenUDP(wg *sync.WaitGroup) {
|
||||||
|
forw.log.Debugf("Start listening on %s:%d", forw.Type, forw.Lport)
|
||||||
|
forw.State = ForwStarted
|
||||||
|
defer wg.Done()
|
||||||
|
stater := func() {
|
||||||
|
forw.State = ForwStopped
|
||||||
|
}
|
||||||
|
defer stater()
|
||||||
|
for {
|
||||||
|
buffer := make([]byte, 2048)
|
||||||
|
size, srcAddr, err := forw.listenUDP.ReadFromUDP(buffer)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Errorf("Error reading: %v", err)
|
||||||
|
continue
|
||||||
|
}
|
||||||
|
go forw.handleUDP(forw.listenUDP, srcAddr, buffer[:size])
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (forw *Forwarder) handleUDP(listConn *net.UDPConn, srcAddr *net.UDPAddr, data []byte) {
|
||||||
|
forw.log.Debugf("Handle on %d started", forw.Lport)
|
||||||
|
if len(forw.Dests) == 0 {
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Select dest address
|
||||||
|
addrnum := rand.Uint32() % uint32(len(forw.Dests))
|
||||||
|
ipaddr := forw.Dests[addrnum].Address
|
||||||
|
destInfo := ipaddr + ":" + strconv.FormatUint(uint64(forw.Dport), 10)
|
||||||
|
|
||||||
|
destAddr, err := net.ResolveUDPAddr("udp", destInfo)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Debugf("Error resolving server address: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
// Write to destination
|
||||||
|
destConn, err := net.DialUDP("udp", nil, destAddr)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Debugf("Error dialing: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
defer destConn.Close()
|
||||||
|
_, err = destConn.Write(data)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Debugf("Error sending message: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
const deadlinePeriod = 5 * time.Second
|
||||||
|
destConn.SetReadDeadline(time.Now().Add(deadlinePeriod))
|
||||||
|
// Read from destination and resend to initiator
|
||||||
|
const readCount = 1
|
||||||
|
for i := 0; i < readCount; i++ {
|
||||||
|
buffer := make([]byte, 1024*2)
|
||||||
|
size, _, err := destConn.ReadFromUDP(buffer)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Debugf("Error reading response: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
_, err = listConn.WriteToUDP(buffer[:size], srcAddr)
|
||||||
|
if err != nil {
|
||||||
|
forw.log.Errorf("Error writing to back: %v", err)
|
||||||
|
return
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
func (forw *Forwarder) Stop() error {
|
func (forw *Forwarder) Stop() error {
|
||||||
return forw.listen.Close()
|
var err error
|
||||||
|
switch forw.Type {
|
||||||
|
case TCP:
|
||||||
|
if forw.listenTCP != nil {
|
||||||
|
return forw.listenTCP.Close()
|
||||||
|
}
|
||||||
|
case UDP:
|
||||||
|
if forw.listenUDP != nil {
|
||||||
|
return forw.listenUDP.Close()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
||||||
@@ -80,23 +198,10 @@ func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
|||||||
if len(forw.Dests) == 0 {
|
if len(forw.Dests) == 0 {
|
||||||
return
|
return
|
||||||
}
|
}
|
||||||
|
// Select dest address
|
||||||
addrnum := rand.Uint32() % uint32(len(forw.Dests))
|
addrnum := rand.Uint32() % uint32(len(forw.Dests))
|
||||||
ipaddr := forw.Dests[addrnum].Address
|
ipaddr := forw.Dests[addrnum].Address
|
||||||
|
|
||||||
/*
|
|
||||||
dstaddr := ipaddr + ":" + strconv.FormatUint(uint64(forw.Dport), 10)
|
|
||||||
outconn, err := net.Dial("tcp", dstaddr)
|
|
||||||
if err != nil {
|
|
||||||
return
|
|
||||||
}
|
|
||||||
defer outconn.Close()
|
|
||||||
var wg sync.WaitGroup
|
|
||||||
wg.Add(1)
|
|
||||||
go forw.stream(&wg, inconn, outconn)
|
|
||||||
wg.Add(1)
|
|
||||||
go forw.stream(&wg, outconn, inconn)
|
|
||||||
wg.Wait()
|
|
||||||
*/
|
|
||||||
str := NewStreamer(forw.ctx)
|
str := NewStreamer(forw.ctx)
|
||||||
err := str.Stream(inconn, ipaddr, forw.Dport)
|
err := str.Stream(inconn, ipaddr, forw.Dport)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
@@ -105,55 +210,6 @@ func (forw *Forwarder) handleTCP(ctx context.Context, inconn net.Conn) {
|
|||||||
forw.log.Debugf("Handler on %d stopped", forw.Lport)
|
forw.log.Debugf("Handler on %d stopped", forw.Lport)
|
||||||
}
|
}
|
||||||
|
|
||||||
func (forw *Forwarder) stream(wg *sync.WaitGroup, inconn io.Reader, outconn io.Writer) {
|
|
||||||
defer wg.Done()
|
|
||||||
_, err := iocopy(forw.ctx, outconn, inconn)
|
|
||||||
if err != nil {
|
|
||||||
forw.log.Errorf("Copy err: %v", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
type Streamer struct {
|
|
||||||
ctx context.Context
|
|
||||||
cancel context.CancelFunc
|
|
||||||
log *logger.Logger
|
|
||||||
}
|
|
||||||
|
|
||||||
func NewStreamer(ctx context.Context) *Streamer {
|
|
||||||
ctx, cancel := context.WithCancel(ctx)
|
|
||||||
log := logger.NewLogger("streamer")
|
|
||||||
return &Streamer{
|
|
||||||
ctx: ctx,
|
|
||||||
cancel: cancel,
|
|
||||||
log: log,
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
func (str *Streamer) Stream(inconn net.Conn, dipaddr string, dport uint32) error {
|
|
||||||
var err error
|
|
||||||
dstaddr := dipaddr + ":" + strconv.FormatUint(uint64(dport), 10)
|
|
||||||
outconn, err := net.Dial("tcp", dstaddr)
|
|
||||||
if err != nil {
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
defer outconn.Close()
|
|
||||||
var wg sync.WaitGroup
|
|
||||||
wg.Add(1)
|
|
||||||
go str.stream(&wg, inconn, outconn)
|
|
||||||
wg.Add(1)
|
|
||||||
go str.stream(&wg, outconn, inconn)
|
|
||||||
wg.Wait()
|
|
||||||
return err
|
|
||||||
}
|
|
||||||
|
|
||||||
func (str *Streamer) stream(wg *sync.WaitGroup, inconn io.Reader, outconn io.Writer) {
|
|
||||||
defer wg.Done()
|
|
||||||
_, err := iocopy(str.ctx, outconn, inconn)
|
|
||||||
if err != nil {
|
|
||||||
str.log.Errorf("Copy err: %v", err)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
type Destination struct {
|
type Destination struct {
|
||||||
Address string `json:"address" yaml:"address"`
|
Address string `json:"address" yaml:"address"`
|
||||||
}
|
}
|
||||||
|
|||||||
+75
-63
@@ -4,7 +4,6 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"fmt"
|
"fmt"
|
||||||
"io"
|
|
||||||
"strings"
|
"strings"
|
||||||
"sync"
|
"sync"
|
||||||
|
|
||||||
@@ -13,10 +12,11 @@ import (
|
|||||||
|
|
||||||
type Proxy struct {
|
type Proxy struct {
|
||||||
Forwarders []*Forwarder `json:"forwarders" yaml:"forwarders"`
|
Forwarders []*Forwarder `json:"forwarders" yaml:"forwarders"`
|
||||||
ctx context.Context `jsreturn errors.New("Zero dport")on:"-" yaml:"-"`
|
ctx context.Context `json:"-" yaml:"-"`
|
||||||
cancel context.CancelFunc `json:"-" yaml:"-"`
|
cancel context.CancelFunc `json:"-" yaml:"-"`
|
||||||
wg sync.WaitGroup `json:"-" yaml:"-"`
|
wg sync.WaitGroup `json:"-" yaml:"-"`
|
||||||
log *logger.Logger `json:"-" yaml:"-"`
|
log *logger.Logger `json:"-" yaml:"-"`
|
||||||
|
mtx sync.Mutex `json:"-" yaml:"-"`
|
||||||
}
|
}
|
||||||
|
|
||||||
func NewProxy() *Proxy {
|
func NewProxy() *Proxy {
|
||||||
@@ -29,106 +29,118 @@ func NewProxy() *Proxy {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
func (bal *Proxy) AddForwarder(ctx context.Context, typ string, lport, dport uint32, addrs ...string) error {
|
func (prox *Proxy) ListForwarders(ctx context.Context) []*Forwarder {
|
||||||
|
prox.mtx.Lock()
|
||||||
|
defer prox.mtx.Unlock()
|
||||||
|
forwards := make([]*Forwarder, 0)
|
||||||
|
for _, forw := range prox.Forwarders {
|
||||||
|
dests := make([]string, 0)
|
||||||
|
for _, dest := range forw.Dests {
|
||||||
|
dests = append(dests, dest.Address)
|
||||||
|
}
|
||||||
|
newForw, _ := NewForwarder(ctx, forw.Type, forw.Lport, forw.Dport, dests...)
|
||||||
|
forwards = append(forwards, newForw)
|
||||||
|
}
|
||||||
|
return forwards
|
||||||
|
}
|
||||||
|
|
||||||
|
func (prox *Proxy) CreateOrUpdateForwarder(ctx context.Context, proto string, lport, dport uint32, addrs ...string) error {
|
||||||
var err error
|
var err error
|
||||||
if lport == 0 {
|
if lport == 0 {
|
||||||
return errors.New("Zero lport")
|
return errors.New("Zero forwarder lport")
|
||||||
}
|
}
|
||||||
if dport == 0 {
|
if dport == 0 {
|
||||||
return errors.New("Zero dport")
|
return errors.New("Zero forwarder dport")
|
||||||
}
|
}
|
||||||
if typ == "" {
|
if proto == "" {
|
||||||
return errors.New("Empty type")
|
return errors.New("Empty forwarder type")
|
||||||
}
|
}
|
||||||
typ = strings.ToLower(typ)
|
proto = strings.ToLower(proto)
|
||||||
if typ != TCP {
|
if proto != TCP && proto != UDP {
|
||||||
return fmt.Errorf("Unknown type %s", typ)
|
return fmt.Errorf("Unknown forwarder protocol %s", proto)
|
||||||
}
|
}
|
||||||
forw, err := NewForwarder(ctx, typ, lport, dport, addrs...)
|
prox.mtx.Lock()
|
||||||
if err != nil {
|
defer prox.mtx.Unlock()
|
||||||
return err
|
|
||||||
|
var forw *Forwarder
|
||||||
|
for _, iforw := range prox.Forwarders {
|
||||||
|
if iforw.Lport == lport && iforw.Type == proto {
|
||||||
|
forw = iforw
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
switch {
|
||||||
|
case forw == nil:
|
||||||
|
prox.log.Debugf("Create rorwarder %s:%d", proto, lport)
|
||||||
|
forw, err = NewForwarder(ctx, proto, lport, dport, addrs...)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
prox.Forwarders = append(prox.Forwarders, forw)
|
||||||
|
prox.wg.Add(1)
|
||||||
|
go forw.Listen(&prox.wg)
|
||||||
|
default:
|
||||||
|
prox.log.Debugf("Update forwarder %s:%d", proto, lport)
|
||||||
|
forw.Dport = dport
|
||||||
|
dests := make([]*Destination, 0)
|
||||||
|
for _, addr := range addrs {
|
||||||
|
dest := NewDestination(addr)
|
||||||
|
forw.Dests = append(dests, dest)
|
||||||
|
}
|
||||||
|
forw.Dests = dests
|
||||||
}
|
}
|
||||||
bal.Forwarders = append(bal.Forwarders, forw)
|
|
||||||
bal.wg.Add(1)
|
|
||||||
go forw.ListenTCP(&bal.wg)
|
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (bal *Proxy) DeleteForwarder(ctx context.Context, lport uint32) error {
|
func (prox *Proxy) DeleteForwarder(ctx context.Context, proto string, lport uint32) error {
|
||||||
var err error
|
var err error
|
||||||
var forw *Forwarder
|
var forw *Forwarder
|
||||||
for _, iforw := range bal.Forwarders {
|
prox.mtx.Lock()
|
||||||
if iforw.Lport == lport {
|
defer prox.mtx.Unlock()
|
||||||
|
for _, iforw := range prox.Forwarders {
|
||||||
|
if iforw.Lport == lport && iforw.Type == proto {
|
||||||
forw = iforw
|
forw = iforw
|
||||||
break
|
break
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
if forw == nil {
|
if forw == nil {
|
||||||
bal.log.Debugf("Forwarder for lport %d not found", lport)
|
prox.log.Debugf("Forwarder for %s:%d not found", proto, lport)
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
bal.log.Debugf("Stop forwarder for lport %d", lport)
|
prox.log.Debugf("Stop forwarder for %s:%d", proto, lport)
|
||||||
err = forw.Stop()
|
err = forw.Stop()
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
forwarders := make([]*Forwarder, 0)
|
forwarders := make([]*Forwarder, 0)
|
||||||
for _, forw := range bal.Forwarders {
|
for _, forw := range prox.Forwarders {
|
||||||
if forw.Lport == lport {
|
if forw.Lport == lport && forw.Type == proto {
|
||||||
continue
|
continue
|
||||||
}
|
}
|
||||||
forwarders = append(forwarders, forw)
|
forwarders = append(forwarders, forw)
|
||||||
}
|
}
|
||||||
bal.Forwarders = forwarders
|
prox.Forwarders = forwarders
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (bal *Proxy) Start() error {
|
func (prox *Proxy) Start() error {
|
||||||
var err error
|
var err error
|
||||||
for _, forw := range bal.Forwarders {
|
prox.mtx.Lock()
|
||||||
bal.wg.Add(1)
|
defer prox.mtx.Unlock()
|
||||||
go forw.ListenTCP(&bal.wg)
|
for _, forw := range prox.Forwarders {
|
||||||
|
prox.wg.Add(1)
|
||||||
|
go forw.Listen(&prox.wg)
|
||||||
}
|
}
|
||||||
bal.wg.Wait()
|
prox.wg.Wait()
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func (bal *Proxy) Stop() error {
|
func (prox *Proxy) Stop() error {
|
||||||
var err error
|
var err error
|
||||||
for _, forw := range bal.Forwarders {
|
prox.mtx.Lock()
|
||||||
|
defer prox.mtx.Unlock()
|
||||||
|
for _, forw := range prox.Forwarders {
|
||||||
forw.Stop()
|
forw.Stop()
|
||||||
}
|
}
|
||||||
return err
|
return err
|
||||||
}
|
}
|
||||||
|
|
||||||
func iocopy(ctx context.Context, writer io.Writer, reader io.Reader) (int64, error) {
|
|
||||||
var err error
|
|
||||||
var size int64
|
|
||||||
var halt bool
|
|
||||||
buffer := make([]byte, 1024*4)
|
|
||||||
for {
|
|
||||||
select {
|
|
||||||
case <-ctx.Done():
|
|
||||||
err = errors.New("Break copy by context")
|
|
||||||
break
|
|
||||||
default:
|
|
||||||
}
|
|
||||||
rsize, err := reader.Read(buffer)
|
|
||||||
if err == io.EOF {
|
|
||||||
err = nil
|
|
||||||
halt = true
|
|
||||||
}
|
|
||||||
if err != nil {
|
|
||||||
return size, err
|
|
||||||
}
|
|
||||||
wsize, err := writer.Write(buffer[0:rsize])
|
|
||||||
size += int64(wsize)
|
|
||||||
if err != nil {
|
|
||||||
return size, err
|
|
||||||
}
|
|
||||||
if halt {
|
|
||||||
break
|
|
||||||
}
|
|
||||||
}
|
|
||||||
return size, err
|
|
||||||
}
|
|
||||||
|
|||||||
@@ -0,0 +1,85 @@
|
|||||||
|
package rproxy
|
||||||
|
|
||||||
|
import (
|
||||||
|
"context"
|
||||||
|
"errors"
|
||||||
|
"io"
|
||||||
|
"net"
|
||||||
|
"strconv"
|
||||||
|
"sync"
|
||||||
|
|
||||||
|
"helmet/app/logger"
|
||||||
|
)
|
||||||
|
|
||||||
|
type Streamer struct {
|
||||||
|
ctx context.Context
|
||||||
|
cancel context.CancelFunc
|
||||||
|
log *logger.Logger
|
||||||
|
}
|
||||||
|
|
||||||
|
func NewStreamer(ctx context.Context) *Streamer {
|
||||||
|
ctx, cancel := context.WithCancel(ctx)
|
||||||
|
log := logger.NewLogger("streamer")
|
||||||
|
return &Streamer{
|
||||||
|
ctx: ctx,
|
||||||
|
cancel: cancel,
|
||||||
|
log: log,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func (str *Streamer) Stream(inconn net.Conn, dipaddr string, dport uint32) error {
|
||||||
|
var err error
|
||||||
|
dstaddr := dipaddr + ":" + strconv.FormatUint(uint64(dport), 10)
|
||||||
|
outconn, err := net.Dial("tcp", dstaddr)
|
||||||
|
if err != nil {
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
defer outconn.Close()
|
||||||
|
var wg sync.WaitGroup
|
||||||
|
wg.Add(1)
|
||||||
|
go str.stream(&wg, inconn, outconn)
|
||||||
|
wg.Add(1)
|
||||||
|
go str.stream(&wg, outconn, inconn)
|
||||||
|
wg.Wait()
|
||||||
|
return err
|
||||||
|
}
|
||||||
|
|
||||||
|
func (str *Streamer) stream(wg *sync.WaitGroup, inconn io.Reader, outconn io.Writer) {
|
||||||
|
defer wg.Done()
|
||||||
|
_, err := iocopy(str.ctx, outconn, inconn)
|
||||||
|
if err != nil {
|
||||||
|
str.log.Errorf("Copy err: %v", err)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
func iocopy(ctx context.Context, writer io.Writer, reader io.Reader) (int64, error) {
|
||||||
|
var err error
|
||||||
|
var size int64
|
||||||
|
var halt bool
|
||||||
|
buffer := make([]byte, 1024*4)
|
||||||
|
for {
|
||||||
|
select {
|
||||||
|
case <-ctx.Done():
|
||||||
|
err = errors.New("Break copy by context")
|
||||||
|
break
|
||||||
|
default:
|
||||||
|
}
|
||||||
|
rsize, err := reader.Read(buffer)
|
||||||
|
if err == io.EOF {
|
||||||
|
err = nil
|
||||||
|
halt = true
|
||||||
|
}
|
||||||
|
if err != nil {
|
||||||
|
return size, err
|
||||||
|
}
|
||||||
|
wsize, err := writer.Write(buffer[0:rsize])
|
||||||
|
size += int64(wsize)
|
||||||
|
if err != nil {
|
||||||
|
return size, err
|
||||||
|
}
|
||||||
|
if halt {
|
||||||
|
break
|
||||||
|
}
|
||||||
|
}
|
||||||
|
return size, err
|
||||||
|
}
|
||||||
@@ -109,13 +109,12 @@ func (svc *Service) logInterceptor(ctx context.Context, req any, info *grpc.Unar
|
|||||||
if err == nil {
|
if err == nil {
|
||||||
svc.log.Debugf("Request: %s", string(reqData))
|
svc.log.Debugf("Request: %s", string(reqData))
|
||||||
}
|
}
|
||||||
res, err := handler(ctx, req)
|
res, handErr := handler(ctx, req)
|
||||||
resData, err := json.Marshal(res)
|
resData, err := json.Marshal(res)
|
||||||
if err == nil {
|
if err == nil {
|
||||||
svc.log.Debugf("Response: %s", string(resData))
|
svc.log.Debugf("Response: %s", string(resData))
|
||||||
}
|
}
|
||||||
|
return res, handErr
|
||||||
return res, err
|
|
||||||
}
|
}
|
||||||
|
|
||||||
func (svc *Service) Stop() {
|
func (svc *Service) Stop() {
|
||||||
|
|||||||
Reference in New Issue
Block a user