Allow parsing of additional net-messages

Author: markus-wa

demoinfocs_test.go

@@ -11,11 +11,13 @@ import (
 	"testing"
 	"time"
 
-	"github.com/markus-wa/godispatch"
+	proto "github.com/gogo/protobuf/proto"
+	dispatch "github.com/markus-wa/godispatch"
 
 	dem "github.com/markus-wa/demoinfocs-golang"
 	common "github.com/markus-wa/demoinfocs-golang/common"
 	events "github.com/markus-wa/demoinfocs-golang/events"
+	msg "github.com/markus-wa/demoinfocs-golang/msg"
 )
 
 const csDemosPath = "test/cs-demos"
@@ -38,6 +40,9 @@ func TestDemoInfoCs(t *testing.T) {
 
 	p := dem.NewParserWithConfig(f, dem.ParserConfig{
 		MsgQueueBufferSize: 1000,
+		AdditionalNetMessageCreators: map[int]dem.NetMessageCreator{
+			4: func() proto.Message { return new(msg.CNETMsg_Tick) },
+		},
 	})
 
 	fmt.Println("Parsing header")
@@ -98,6 +103,13 @@ func TestDemoInfoCs(t *testing.T) {
 		}
 	})
 
+	// Net-message stuff
+	var netTickHandlerID dispatch.HandlerIdentifier
+	netTickHandlerID = p.RegisterNetMessageHandler(func(tick *msg.CNETMsg_Tick) {
+		fmt.Println("Net-message tick handled, unregistering - tick:", tick.Tick)
+		p.UnregisterNetMessageHandler(netTickHandlerID)
+	})
+
 	ts := time.Now()
 	var done int64
 	go func() {

demopacket.go

@@ -65,7 +65,14 @@ func (p *Parser) parsePacket() {
 				debugUnhandledMessage(cmd, name)
 			}
 
-			if name == "" {
+			if name != "" {
+				// Handle additional net-messages as defined by the user
+				creator := p.additionalNetMessageCreators[cmd]
+				if creator != nil {
+					m = creator()
+					break
+				}
+			} else {
 				// Send a warning if the command is unknown
 				// This might mean our proto files are out of date
 				p.eventDispatcher.Dispatch(events.ParserWarnEvent{Message: fmt.Sprintf("Unknown message command %q", cmd)})
@@ -94,6 +101,10 @@ func (p *Parser) parsePacket() {
 	p.bitReader.EndChunk()
 }
 
+// NetMessageCreator creates additional net-messages to be dispatched to net-message handlers.
+// See also: ParserConfig.AdditionalNetMessageCreators & Parser.RegisterNetMessageHandler()
+type NetMessageCreator func() proto.Message
+
 /*
 Format of 'CommandInfos' - I honestly have no clue what they are good for.
 If you find a use for this please let me know!

parser.go

@@ -19,31 +19,32 @@ import (
 // After parsing the header Parser.ParseNextFrame() and Parser.ParseToEnd() can be used to parse the demo.
 // Use Parser.RegisterEventHandler() to receive notifications about events.
 type Parser struct {
-	bitReader             *bit.BitReader
-	stParser              st.SendTableParser
-	msgDispatcher         dp.Dispatcher
-	eventDispatcher       dp.Dispatcher
-	msgQueue              chan interface{}
-	gameState             GameState
-	currentFrame          int
-	bombsiteA             bombsite
-	bombsiteB             bombsite
-	header                *common.DemoHeader // Pointer so we can check for nil
-	equipmentMapping      map[*st.ServerClass]common.EquipmentElement
-	rawPlayers            map[int]*playerInfo
-	entityIDToPlayers     map[int]*common.Player // Temporary storage since we need to map players from entityID to userID later
-	additionalPlayerInfo  [maxPlayers]common.AdditionalPlayerInformation
-	entities              map[int]*st.Entity
-	modelPreCache         []string                      // Used to find out whether a weapon is a p250 or cz for example (same id)
-	weapons               [maxEntities]common.Equipment // Used to remember what a weapon is (p250 / cz etc.)
-	triggers              map[int]*boundingBoxInformation
-	instanceBaselines     map[int][]byte
-	preprocessedBaselines map[int]map[int]st.PropValue
-	gameEventDescs        map[int32]*msg.CSVCMsg_GameEventListDescriptorT
-	stringTables          []*msg.CSVCMsg_CreateStringTable
-	cancelChan            chan struct{}
-	err                   error
-	errLock               sync.Mutex
+	bitReader                    *bit.BitReader
+	stParser                     st.SendTableParser
+	msgDispatcher                dp.Dispatcher
+	additionalNetMessageCreators map[int]NetMessageCreator
+	eventDispatcher              dp.Dispatcher
+	msgQueue                     chan interface{}
+	gameState                    GameState
+	currentFrame                 int
+	bombsiteA                    bombsite
+	bombsiteB                    bombsite
+	header                       *common.DemoHeader // Pointer so we can check for nil
+	equipmentMapping             map[*st.ServerClass]common.EquipmentElement
+	rawPlayers                   map[int]*playerInfo
+	entityIDToPlayers            map[int]*common.Player // Temporary storage since we need to map players from entityID to userID later
+	additionalPlayerInfo         [maxPlayers]common.AdditionalPlayerInformation
+	entities                     map[int]*st.Entity
+	modelPreCache                []string                      // Used to find out whether a weapon is a p250 or cz for example (same id)
+	weapons                      [maxEntities]common.Equipment // Used to remember what a weapon is (p250 / cz etc.)
+	triggers                     map[int]*boundingBoxInformation
+	instanceBaselines            map[int][]byte
+	preprocessedBaselines        map[int]map[int]st.PropValue
+	gameEventDescs               map[int32]*msg.CSVCMsg_GameEventListDescriptorT
+	stringTables                 []*msg.CSVCMsg_CreateStringTable
+	cancelChan                   chan struct{}
+	err                          error
+	errLock                      sync.Mutex
 }
 
 type bombsite struct {
@@ -100,17 +101,33 @@ func (p *Parser) Progress() float32 {
 // Must be of type func(<EventType>) where EventType is the kind of event that is handled.
 // To catch all events func(interface{}) can be used.
 // Parameter handler has to be of type interface{} because lolnogenerics.
-// Returns a identifier with which the handler can be removed via UnregisterEventHandler()
+// Returns a identifier with which the handler can be removed via UnregisterEventHandler().
 func (p *Parser) RegisterEventHandler(handler interface{}) dp.HandlerIdentifier {
 	return p.eventDispatcher.RegisterHandler(handler)
 }
 
-// UnregisterEventHandler removes a handler via identifier.
-// The identifier is returned at registration by RegisterEventHandler()
+// UnregisterEventHandler removes a game event handler via identifier.
+// The identifier is returned at registration by RegisterEventHandler().
 func (p *Parser) UnregisterEventHandler(identifier dp.HandlerIdentifier) {
 	p.eventDispatcher.UnregisterHandler(identifier)
 }
 
+// RegisterNetMessageHandler registers a handler for net-messages.
+// Must be of type func(*<EventType>) where EventType is the kind of event that is handled.
+// To catch all events func(interface{}) can be used.
+// Parameter handler has to be of type interface{} because lolnogenerics.
+// Returns a identifier with which the handler can be removed via UnregisterNetMessageHandler().
+// This is a beta feature and may be changed or replaced without notice.
+func (p *Parser) RegisterNetMessageHandler(handler interface{}) dp.HandlerIdentifier {
+	return p.msgDispatcher.RegisterHandler(handler)
+}
+
+// UnregisterNetMessageHandler removes a net-message handler via identifier.
+// The identifier is returned at registration by RegisterNetMessageHandler().
+func (p *Parser) UnregisterNetMessageHandler(identifier dp.HandlerIdentifier) {
+	p.msgDispatcher.UnregisterHandler(identifier)
+}
+
 func (p *Parser) error() (err error) {
 	p.errLock.Lock()
 	err = p.err
@@ -143,6 +160,12 @@ type ParserConfig struct {
 	// this is the default behavior for DefaultParserConfig.
 	// Zero enforces sequential parsing.
 	MsgQueueBufferSize int
+	// AdditionalNetMessageCreators maps net-message-IDs to creators (instantiators).
+	// The creators should return a new instance of the correct protobuf-message type (from the msg package).
+	// Interesting net-message-IDs can easily be discovered with the build-tag 'debugdemoinfocs'; when looking for 'UnhandledMessage'.
+	// Check out demopacket.go to see which net-messages are already being parsed by default.
+	// This is a beta feature and may be changed or replaced without notice.
+	AdditionalNetMessageCreators map[int]NetMessageCreator
 }
 
 // DefaultParserConfig is the default Parser configuration used by NewParser().
@@ -182,6 +205,8 @@ func NewParserWithConfig(demostream io.Reader, config ParserConfig) *Parser {
 		p.initMsgQueue(config.MsgQueueBufferSize)
 	}
 
+	p.additionalNetMessageCreators = config.AdditionalNetMessageCreators
+
 	return &p
 }