hoelee/cloudflared
1
0
Fork 0
mirror of https://github.com/cloudflare/cloudflared.git synced 2025-07-27 21:20:00 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

TUN-6688: Update RegisterUdpSession capnproto to include trace context

Browse Source
This commit is contained in:
Devin Carr
2022-09-07 15:06:06 -07:00
parent 11cbff4ff7
commit e380333520
8 changed files with 361 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tunnelrpc/pogs/sessionrpc.go
View File

@@ -14,7 +14,7 @@ import (
)
type SessionManager interface {
RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeAfterIdleHint time.Duration) error
RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeAfterIdleHint time.Duration, traceContext string) error
UnregisterUdpSession(ctx context.Context, sessionID uuid.UUID, message string) error
}
@@ -50,12 +50,19 @@ func (i SessionManager_PogsImpl) RegisterUdpSession(p tunnelrpc.SessionManager_r
closeIdleAfterHint := time.Duration(p.Params.CloseAfterIdleHint())
traceContext, err := p.Params.TraceContext()
if err != nil {
return err
}
resp := RegisterUdpSessionResponse{}
registrationErr := i.impl.RegisterUdpSession(p.Ctx, sessionID, dstIP, dstPort, closeIdleAfterHint)
registrationErr := i.impl.RegisterUdpSession(p.Ctx, sessionID, dstIP, dstPort, closeIdleAfterHint, traceContext)
if registrationErr != nil {
resp.Err = registrationErr
}
// TUN-6689: Add spans to return path for RegisterUdpSession
result, err := p.Results.NewResult()
if err != nil {
return err
@@ -85,13 +92,17 @@ func (i SessionManager_PogsImpl) UnregisterUdpSession(p tunnelrpc.SessionManager
}
type RegisterUdpSessionResponse struct {
Err error
Err error
Spans []byte // Spans in protobuf format
}
func (p *RegisterUdpSessionResponse) Marshal(s tunnelrpc.RegisterUdpSessionResponse) error {
if p.Err != nil {
return s.SetErr(p.Err.Error())
}
if err := s.SetSpans(p.Spans); err != nil {
return err
}
return nil
}
@@ -103,6 +114,10 @@ func (p *RegisterUdpSessionResponse) Unmarshal(s tunnelrpc.RegisterUdpSessionRes
if respErr != "" {
p.Err = fmt.Errorf(respErr)
}
p.Spans, err = s.Spans()
if err != nil {
return err
}
return nil
}
@@ -116,7 +131,7 @@ func (c SessionManager_PogsClient) Close() error {
return c.Conn.Close()
}
func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeAfterIdleHint time.Duration) (*RegisterUdpSessionResponse, error) {
func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessionID uuid.UUID, dstIP net.IP, dstPort uint16, closeAfterIdleHint time.Duration, traceContext string) (*RegisterUdpSessionResponse, error) {
client := tunnelrpc.SessionManager{Client: c.Client}
promise := client.RegisterUdpSession(ctx, func(p tunnelrpc.SessionManager_registerUdpSession_Params) error {
if err := p.SetSessionId(sessionID[:]); err != nil {
@@ -127,6 +142,7 @@ func (c SessionManager_PogsClient) RegisterUdpSession(ctx context.Context, sessi
}
p.SetDstPort(dstPort)
p.SetCloseAfterIdleHint(int64(closeAfterIdleHint))
p.SetTraceContext(traceContext)
return nil
})
result, err := promise.Result().Struct()