-
Notifications
You must be signed in to change notification settings - Fork 0
/
adapter.go
89 lines (80 loc) · 2.41 KB
/
adapter.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
// Package adapter is an adapter for the hrpc protocol.
package adapter
import (
"context"
"os"
"github.com/gofiber/fiber/v2"
"github.com/gofiber/websocket/v2"
goserver "github.com/harmony-development/hrpc/pkg/go-server"
"github.com/harmony-development/legato/adapter/hrpc"
"github.com/harmony-development/legato/api/middleware"
"github.com/harmony-development/legato/util/panic_fmt"
)
func RegisterHandler[T context.Context](mid *middleware.MiddlewareHandler, router fiber.Router, handler goserver.MethodHandler[T]) {
for route, h := range handler.Routes() {
router.All(route, mid.Auth(route), UnaryHandler(h))
}
for route, h := range handler.StreamRoutes() {
router.All(route, mid.Auth(route), StreamHandler(h))
}
}
func UnaryHandler[T context.Context](h goserver.UnaryMethodData[T]) fiber.Handler {
return func(c *fiber.Ctx) error {
input, err := hrpc.UnmarshalHRPC(c.Body(), c.Get("Content-Type"), h.Input)
if err != nil {
return err
}
result, err := h.Handler(c.UserContext().(T), input)
if err != nil {
return err
}
response, err := hrpc.MarshalHRPC(result, c.Get("Content-Type"))
return c.Send(response)
}
}
func StreamHandler[T context.Context](h goserver.StreamMethodData[T]) fiber.Handler {
return func(c *fiber.Ctx) error {
ctx := c.UserContext()
return websocket.New(func(conn *websocket.Conn) {
reqChan := make(chan goserver.VTProtoMessage)
sendChan := make(chan goserver.VTProtoMessage)
go func() {
defer func() {
close(reqChan)
close(sendChan)
if r := recover(); r != nil {
panic_fmt.RecoverError(os.Getenv("RAW_LOG") != "true", r)
}
}()
h.Handler(ctx.(T), reqChan, sendChan)
}()
go func() {
defer func() {
if r := recover(); r != nil {
panic_fmt.RecoverError(os.Getenv("RAW_LOG") != "true", r)
}
}()
_, msg, err := conn.ReadMessage() // Read the next message from the client.
if err != nil {
return
}
hrpcMsg, err := hrpc.UnmarshalHRPC(msg, "application/hrpc", h.Input)
if err != nil {
conn.Close() // TODO: don't close socket, send error to client
return
}
reqChan <- hrpcMsg
}()
for msg := range sendChan {
resp, err := hrpc.MarshalHRPC(msg, "application/hrpc")
if err != nil {
conn.Close()
return
}
conn.WriteMessage(websocket.BinaryMessage, append([]byte{0}, resp...))
}
}, websocket.Config{
Subprotocols: []string{"hrpc1"},
})(c)
}
}