chore: code cleanup

This commit is contained in:
wwqgtxx
2023-10-11 22:54:19 +08:00
parent 0dc6a726c1
commit 129283066f
14 changed files with 97 additions and 106 deletions

View File

@ -17,35 +17,15 @@ func WithAdditions(ctx context.Context, additions ...inbound.Addition) context.C
return context.WithValue(ctx, ctxKeyAdditions, additions)
}
func getAdditions(ctx context.Context) []inbound.Addition {
func getAdditions(ctx context.Context) (additions []inbound.Addition) {
if v := ctx.Value(ctxKeyAdditions); v != nil {
if a, ok := v.([]inbound.Addition); ok {
return a
additions = a
}
}
return nil
}
func combineAdditions(ctx context.Context, additions []inbound.Addition, extraAdditions ...inbound.Addition) []inbound.Addition {
additionsCloned := false
if ctxAdditions := getAdditions(ctx); len(ctxAdditions) > 0 {
additions = slices.Clone(additions)
additionsCloned = true
additions = append(additions, ctxAdditions...)
}
if user, ok := auth.UserFromContext[string](ctx); ok {
if !additionsCloned {
additions = slices.Clone(additions)
additionsCloned = true
}
additions = slices.Clone(additions)
additions = append(additions, inbound.WithInUser(user))
}
if len(extraAdditions) > 0 {
if !additionsCloned {
additions = slices.Clone(additions)
additionsCloned = true
}
additions = append(additions, extraAdditions...)
}
return additions
return
}

View File

@ -92,12 +92,10 @@ func (h *ListenerHandler) NewConnection(ctx context.Context, conn net.Conn, meta
cMetadata := &C.Metadata{
NetWork: C.TCP,
Type: h.Type,
Host: metadata.Destination.Fqdn,
}
additions := combineAdditions(ctx, h.Additions, inbound.WithDstAddr(metadata.Destination), inbound.WithSrcAddr(metadata.Source), inbound.WithInAddr(conn.LocalAddr()))
for _, addition := range additions {
addition.Apply(cMetadata)
}
inbound.ApplyAdditions(cMetadata, inbound.WithDstAddr(metadata.Destination), inbound.WithSrcAddr(metadata.Source), inbound.WithInAddr(conn.LocalAddr()))
inbound.ApplyAdditions(cMetadata, getAdditions(ctx)...)
inbound.ApplyAdditions(cMetadata, h.Additions...)
h.Tunnel.HandleTCPConn(conn, cMetadata) // this goroutine must exit after conn unused
return nil
@ -155,12 +153,10 @@ func (h *ListenerHandler) NewPacketConnection(ctx context.Context, conn network.
cMetadata := &C.Metadata{
NetWork: C.UDP,
Type: h.Type,
Host: dest.Fqdn,
}
additions := combineAdditions(ctx, h.Additions, inbound.WithDstAddr(dest), inbound.WithSrcAddr(metadata.Source), inbound.WithInAddr(conn.LocalAddr()))
for _, addition := range additions {
addition.Apply(cMetadata)
}
inbound.ApplyAdditions(cMetadata, inbound.WithDstAddr(dest), inbound.WithSrcAddr(metadata.Source), inbound.WithInAddr(conn.LocalAddr()))
inbound.ApplyAdditions(cMetadata, getAdditions(ctx)...)
inbound.ApplyAdditions(cMetadata, h.Additions...)
h.Tunnel.HandleUDPPacket(cPacket, cMetadata)
}