95 lines
2.3 KiB
Go
95 lines
2.3 KiB
Go
package grpcserver
|
|
|
|
import (
|
|
"log"
|
|
"strings"
|
|
|
|
"relay/device"
|
|
"relay/registry"
|
|
|
|
tunnelpb "relay/proto/tunnel"
|
|
)
|
|
|
|
type TunnelServiceImpl struct {
|
|
tunnelpb.UnimplementedTunnelServiceServer
|
|
}
|
|
|
|
func (s *TunnelServiceImpl) Tunnel(stream tunnelpb.TunnelService_TunnelServer) error {
|
|
dev := device.NewDevice(stream)
|
|
log.Println("Device connected")
|
|
|
|
defer func() {
|
|
log.Println("Device disconnected")
|
|
dev.Close()
|
|
registry.Global.Mu.Lock()
|
|
for d, v := range registry.Global.Domains {
|
|
if v == dev {
|
|
registry.Global.Domains[d] = nil
|
|
log.Println("Domain unbound from device:", d)
|
|
}
|
|
}
|
|
registry.Global.Mu.Unlock()
|
|
}()
|
|
|
|
for {
|
|
frame, err := stream.Recv()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
switch frame.Type {
|
|
|
|
// case tunnelpb.FrameType_FRAME_OPEN:
|
|
// domain := strings.ToLower(string(frame.Payload))
|
|
// log.Println("BIND request from device for domain:", domain)
|
|
// registry.Global.Mu.Lock()
|
|
// if _, ok := registry.Global.Domains[domain]; ok {
|
|
// registry.Global.Domains[domain] = dev
|
|
// log.Println("Domain bound to device:", domain)
|
|
// } else {
|
|
// log.Println("Bind rejected, domain not registered:", domain)
|
|
// }
|
|
// registry.Global.Mu.Unlock()
|
|
|
|
case tunnelpb.FrameType_FRAME_BIND_REQUEST:
|
|
domain := strings.ToLower(string(frame.Payload))
|
|
log.Println("BIND request from device for domain:", domain)
|
|
|
|
registry.Global.Mu.Lock()
|
|
if _, ok := registry.Global.Domains[domain]; ok {
|
|
registry.Global.Domains[domain] = dev
|
|
registry.Global.Mu.Unlock()
|
|
|
|
log.Println("Domain bound to device:", domain)
|
|
dev.SendFrame(&tunnelpb.Frame{
|
|
Type: tunnelpb.FrameType_FRAME_BIND_OK,
|
|
Payload: []byte(domain),
|
|
})
|
|
} else {
|
|
registry.Global.Mu.Unlock()
|
|
|
|
log.Println("Bind rejected, domain not registered:", domain)
|
|
dev.SendFrame(&tunnelpb.Frame{
|
|
Type: tunnelpb.FrameType_FRAME_BIND_REJECTED,
|
|
Payload: []byte(domain),
|
|
})
|
|
}
|
|
|
|
case tunnelpb.FrameType_FRAME_DATA:
|
|
if c, ok := dev.GetClient(frame.StreamId); ok {
|
|
if _, err := c.Write(frame.Payload); err != nil {
|
|
dev.RemoveStream(frame.StreamId)
|
|
}
|
|
}
|
|
|
|
case tunnelpb.FrameType_FRAME_CLOSE:
|
|
dev.RemoveStream(frame.StreamId)
|
|
|
|
case tunnelpb.FrameType_FRAME_PING:
|
|
dev.SendFrame(&tunnelpb.Frame{
|
|
Type: tunnelpb.FrameType_FRAME_PONG,
|
|
})
|
|
}
|
|
}
|
|
}
|