fix: cleanup contextToRootID map on proxy done and map name correction (#92)

Co-authored-by: Wim Spaargaren <[email protected]>

Author: wimspaargaren

proxywasm/abi_lifecycle.go

@@ -29,6 +29,9 @@ func proxyOnContextCreate(contextID uint32, rootContextID uint32) {
 
 //export proxy_on_done
 func proxyOnDone(contextID uint32) bool {
+	defer func() {
+		delete(currentState.contextIDToRootID, contextID)
+	}()
 	if ctx, ok := currentState.streams[contextID]; ok {
 		currentState.setActiveContextID(contextID)
 		delete(currentState.streams, contextID)

proxywasm/abi_lifecycle_test.go

@@ -14,10 +14,10 @@ func Test_proxyOnContextCreate(t *testing.T) {
 
 	var cnt int
 	currentState = &state{
-		rootContexts:     map[uint32]*rootContextState{},
-		httpStreams:      map[uint32]HttpContext{},
-		streams:          map[uint32]StreamContext{},
-		contextIDToRooID: map[uint32]uint32{},
+		rootContexts:      map[uint32]*rootContextState{},
+		httpStreams:       map[uint32]HttpContext{},
+		streams:           map[uint32]StreamContext{},
+		contextIDToRootID: map[uint32]uint32{},
 	}
 
 	SetNewRootContext(func(contextID uint32) RootContext {

proxywasm/vmstate.go

@@ -34,15 +34,15 @@ type state struct {
 	newHttpContext   func(rootContextID, contextID uint32) HttpContext
 	httpStreams      map[uint32]HttpContext
 
-	contextIDToRooID map[uint32]uint32
-	activeContextID  uint32
+	contextIDToRootID map[uint32]uint32
+	activeContextID   uint32
 }
 
 var currentState = &state{
-	rootContexts:     make(map[uint32]*rootContextState),
-	httpStreams:      make(map[uint32]HttpContext),
-	streams:          make(map[uint32]StreamContext),
-	contextIDToRooID: make(map[uint32]uint32),
+	rootContexts:      make(map[uint32]*rootContextState),
+	httpStreams:       make(map[uint32]HttpContext),
+	streams:           make(map[uint32]StreamContext),
+	contextIDToRootID: make(map[uint32]uint32),
 }
 
 func SetNewRootContext(f func(contextID uint32) RootContext) {
@@ -85,7 +85,7 @@ func (s *state) createStreamContext(contextID uint32, rootContextID uint32) {
 	}
 
 	ctx := s.newStreamContext(rootContextID, contextID)
-	s.contextIDToRooID[contextID] = rootContextID
+	s.contextIDToRootID[contextID] = rootContextID
 	s.streams[contextID] = ctx
 }
 
@@ -99,12 +99,12 @@ func (s *state) createHttpContext(contextID uint32, rootContextID uint32) {
 	}
 
 	ctx := s.newHttpContext(rootContextID, contextID)
-	s.contextIDToRooID[contextID] = rootContextID
+	s.contextIDToRootID[contextID] = rootContextID
 	s.httpStreams[contextID] = ctx
 }
 
 func (s *state) registerHttpCallOut(calloutID uint32, callback HttpCalloutCallBack) {
-	r := s.rootContexts[s.contextIDToRooID[s.activeContextID]]
+	r := s.rootContexts[s.contextIDToRootID[s.activeContextID]]
 	r.httpCallbacks[calloutID] = &struct {
 		callback        HttpCalloutCallBack
 		callerContextID uint32

proxywasm/vmstate_test.go

@@ -84,10 +84,10 @@ func TestState_createStreamContext(t *testing.T) {
 		rid uint32 = 10
 	)
 	s := &state{
-		rootContexts:     map[uint32]*rootContextState{rid: nil},
-		streams:          map[uint32]StreamContext{},
-		newStreamContext: func(rootContextID, contextID uint32) StreamContext { return &sc{} },
-		contextIDToRooID: map[uint32]uint32{},
+		rootContexts:      map[uint32]*rootContextState{rid: nil},
+		streams:           map[uint32]StreamContext{},
+		newStreamContext:  func(rootContextID, contextID uint32) StreamContext { return &sc{} },
+		contextIDToRootID: map[uint32]uint32{},
 	}
 
 	s.createStreamContext(cid, rid)
@@ -105,10 +105,10 @@ func TestState_createHttpContext(t *testing.T) {
 		rid uint32 = 10
 	)
 	s := &state{
-		rootContexts:     map[uint32]*rootContextState{rid: nil},
-		httpStreams:      map[uint32]HttpContext{},
-		newHttpContext:   func(rootContextID, contextID uint32) HttpContext { return &hc{} },
-		contextIDToRooID: map[uint32]uint32{},
+		rootContexts:      map[uint32]*rootContextState{rid: nil},
+		httpStreams:       map[uint32]HttpContext{},
+		newHttpContext:    func(rootContextID, contextID uint32) HttpContext { return &hc{} },
+		contextIDToRootID: map[uint32]uint32{},
 	}
 
 	s.createHttpContext(cid, rid)

proxywasm/vmstate_test_export.go

@@ -19,10 +19,10 @@ package proxywasm
 func VMStateReset() {
 	// (@mathetake) I assume that the currentState be protected by lock on hostMux
 	currentState = &state{
-		rootContexts:     make(map[uint32]*rootContextState),
-		httpStreams:      make(map[uint32]HttpContext),
-		streams:          make(map[uint32]StreamContext),
-		contextIDToRooID: make(map[uint32]uint32),
+		rootContexts:      make(map[uint32]*rootContextState),
+		httpStreams:       make(map[uint32]HttpContext),
+		streams:           make(map[uint32]StreamContext),
+		contextIDToRootID: make(map[uint32]uint32),
 	}
 }