Merge pull request optimism-java#141 from fearlessfe/beacon

feat: refactor protocolId
fearlessfe · Aug 29, 2024 · 1433d80 · 1433d80
2 parents 75885b5 + d4d644c
commit 1433d80
Show file tree

Hide file tree

Showing 7 changed files with 56 additions and 32 deletions.
diff --git a/cmd/shisui/main.go b/cmd/shisui/main.go
@@ -121,14 +121,14 @@ func startPortalRpcServer(config Config, conn discover.UDPConn, addr string) err
 		return err
 	}
 
-	if slices.Contains(config.Networks, portalwire.HistoryNetworkName) {
+	if slices.Contains(config.Networks, portalwire.History.Name()) {
 		err = initHistory(config, server, conn, localNode, discV5)
 		if err != nil {
 			return err
 		}
 	}
 
-	if slices.Contains(config.Networks, portalwire.BeaconNetworkName) {
+	if slices.Contains(config.Networks, portalwire.Beacon.Name()) {
 		err = initBeacon(config, server, conn, localNode, discV5)
 		if err != nil {
 			return err
@@ -181,7 +181,7 @@ func initHistory(config Config, server *rpc.Server, conn discover.UDPConn, local
 	}
 	contentQueue := make(chan *discover.ContentElement, 50)
 
-	protocol, err := discover.NewPortalProtocol(config.Protocol, string(portalwire.HistoryNetwork), config.PrivateKey, conn, localNode, discV5, contentStorage, contentQueue)
+	protocol, err := discover.NewPortalProtocol(config.Protocol, portalwire.History, config.PrivateKey, conn, localNode, discV5, contentStorage, contentQueue)
 
 	if err != nil {
 		return err
@@ -222,7 +222,7 @@ func initBeacon(config Config, server *rpc.Server, conn discover.UDPConn, localN
 	}
 	contentQueue := make(chan *discover.ContentElement, 50)
 
-	protocol, err := discover.NewPortalProtocol(config.Protocol, string(portalwire.BeaconLightClientNetwork), config.PrivateKey, conn, localNode, discV5, contentStorage, contentQueue)
+	protocol, err := discover.NewPortalProtocol(config.Protocol, portalwire.Beacon, config.PrivateKey, conn, localNode, discV5, contentStorage, contentQueue)
 
 	if err != nil {
 		return err

diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
@@ -1033,7 +1033,7 @@ Please note that --` + MetricsHTTPFlag.Name + ` must be set to start the server.
 		Name:     "networks",
 		Usage:    "Portal sub networks: history, beacon, state",
 		Category: flags.PortalNetworkCategory,
-		Value:    cli.NewStringSlice(portalwire.HistoryNetworkName),
+		Value:    cli.NewStringSlice(portalwire.History.Name()),
 	}
 )
 

diff --git a/p2p/discover/portal_protocol.go b/p2p/discover/portal_protocol.go
@@ -206,16 +206,15 @@ func defaultContentIdFunc(contentKey []byte) []byte {
 	return digest[:]
 }
 
-func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateKey *ecdsa.PrivateKey, conn UDPConn, localNode *enode.LocalNode, discV5 *UDPv5, storage storage.ContentStorage, contentQueue chan *ContentElement, opts ...PortalProtocolOption) (*PortalProtocol, error) {
+func NewPortalProtocol(config *PortalProtocolConfig, protocolId portalwire.ProtocolId, privateKey *ecdsa.PrivateKey, conn UDPConn, localNode *enode.LocalNode, discV5 *UDPv5, storage storage.ContentStorage, contentQueue chan *ContentElement, opts ...PortalProtocolOption) (*PortalProtocol, error) {
 	closeCtx, cancelCloseCtx := context.WithCancel(context.Background())
 
-	protocolName := portalwire.NetworkNameMap[protocolId]
 	protocol := &PortalProtocol{
 		cachedIds:      make(map[string]enode.ID),
-		protocolId:     protocolId,
-		protocolName:   protocolName,
+		protocolId:     string(protocolId),
+		protocolName:   protocolId.Name(),
 		ListenAddr:     config.ListenAddr,
-		Log:            log.New("protocol", protocolName),
+		Log:            log.New("protocol", protocolId.Name()),
 		PrivateKey:     privateKey,
 		NetRestrict:    config.NetRestrict,
 		BootstrapNodes: config.BootstrapNodes,
@@ -250,7 +249,7 @@ func (p *PortalProtocol) Start() error {
 	}
 
 	p.DiscV5.RegisterTalkHandler(p.protocolId, p.handleTalkRequest)
-	p.DiscV5.RegisterTalkHandler(string(portalwire.UTPNetwork), p.handleUtpTalkRequest)
+	p.DiscV5.RegisterTalkHandler(string(portalwire.Utp), p.handleUtpTalkRequest)
 
 	go p.table.loop()
 
@@ -326,7 +325,7 @@ func (p *PortalProtocol) setupUDPListening() error {
 			defer p.cachedIdsLock.Unlock()
 			if id, ok := p.cachedIds[addr.String()]; ok {
 				//_, err := p.DiscV5.TalkRequestToID(id, addr, string(portalwire.UTPNetwork), buf)
-				req := &v5wire.TalkRequest{Protocol: string(portalwire.UTPNetwork), Message: buf}
+				req := &v5wire.TalkRequest{Protocol: string(portalwire.Utp), Message: buf}
 				p.DiscV5.sendFromAnotherThread(id, netip.AddrPortFrom(netutil.IPToAddr(addr.IP), uint16(addr.Port)), req)
 
 				return len(buf), err

diff --git a/p2p/discover/portal_protocol_test.go b/p2p/discover/portal_protocol_test.go
@@ -85,7 +85,7 @@ func setupLocalPortalNode(addr string, bootNodes []*enode.Node) (*PortalProtocol
 	}
 
 	contentQueue := make(chan *ContentElement, 50)
-	portalProtocol, err := NewPortalProtocol(conf, string(portalwire.HistoryNetwork), privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
+	portalProtocol, err := NewPortalProtocol(conf, portalwire.History, privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
 	if err != nil {
 		return nil, err
 	}

diff --git a/p2p/discover/portalwire/messages.go b/p2p/discover/portalwire/messages.go
@@ -37,29 +37,54 @@ const (
 )
 
 // Protocol IDs for the portal protocol.
+// var (
+// 	StateNetwork             = []byte{0x50, 0x0a}
+// 	HistoryNetwork           = []byte{0x50, 0x0b}
+// 	TxGossipNetwork          = []byte{0x50, 0x0c}
+// 	HeaderGossipNetwork      = []byte{0x50, 0x0d}
+// 	CanonicalIndicesNetwork  = []byte{0x50, 0x0e}
+// 	BeaconLightClientNetwork = []byte{0x50, 0x1a}
+// 	UTPNetwork               = []byte{0x75, 0x74, 0x70}
+// 	Rendezvous               = []byte{0x72, 0x65, 0x6e}
+// )
+
+type ProtocolId []byte
+
 var (
-	StateNetwork             = []byte{0x50, 0x0a}
-	HistoryNetwork           = []byte{0x50, 0x0b}
-	TxGossipNetwork          = []byte{0x50, 0x0c}
-	HeaderGossipNetwork      = []byte{0x50, 0x0d}
-	CanonicalIndicesNetwork  = []byte{0x50, 0x0e}
-	BeaconLightClientNetwork = []byte{0x50, 0x1a}
-	UTPNetwork               = []byte{0x75, 0x74, 0x70}
-	Rendezvous               = []byte{0x72, 0x65, 0x6e}
+	State             ProtocolId = []byte{0x50, 0x0A}
+	History           ProtocolId = []byte{0x50, 0x0B}
+	Beacon            ProtocolId = []byte{0x50, 0x0C}
+	CanonicalIndices  ProtocolId = []byte{0x50, 0x0D}
+	VerkleState       ProtocolId = []byte{0x50, 0x0E}
+	TransactionGossip ProtocolId = []byte{0x50, 0x0F}
+	Utp               ProtocolId = []byte{0x75, 0x74, 0x70}
 )
 
-const (
-	HistoryNetworkName = "history"
-	BeaconNetworkName  = "beacon"
-	StateNetworkName   = "state"
-)
+var protocalName = map[string]string{
+	string(State):             "state",
+	string(History):           "history",
+	string(Beacon):            "beacon",
+	string(CanonicalIndices):  "canonical indices",
+	string(VerkleState):       "verkle state",
+	string(TransactionGossip): "transaction gossip",
+}
 
-var NetworkNameMap = map[string]string{
-	string(StateNetwork):             StateNetworkName,
-	string(HistoryNetwork):           HistoryNetworkName,
-	string(BeaconLightClientNetwork): BeaconNetworkName,
+func (p ProtocolId) Name() string {
+	return protocalName[string(p)]
 }
 
+// const (
+// 	HistoryNetworkName = "history"
+// 	BeaconNetworkName  = "beacon"
+// 	StateNetworkName   = "state"
+// )
+
+// var NetworkNameMap = map[string]string{
+// 	string(StateNetwork):             StateNetworkName,
+// 	string(HistoryNetwork):           HistoryNetworkName,
+// 	string(BeaconLightClientNetwork): BeaconNetworkName,
+// }
+
 type ContentKV struct {
 	ContentKey []byte
 	Content    []byte

diff --git a/portalnetwork/beacon/beacon_network_test.go b/portalnetwork/beacon/beacon_network_test.go
@@ -65,7 +65,7 @@ func setupBeaconNetwork(addr string, bootNodes []*enode.Node) (*BeaconNetwork, e
 
 	contentQueue := make(chan *discover.ContentElement, 50)
 
-	portalProtocol, err := discover.NewPortalProtocol(conf, string(portalwire.BeaconLightClientNetwork), privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
+	portalProtocol, err := discover.NewPortalProtocol(conf, portalwire.Beacon, privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
 	if err != nil {
 		return nil, err
 	}

diff --git a/portalnetwork/history/history_network_test.go b/portalnetwork/history/history_network_test.go
@@ -419,7 +419,7 @@ func genHistoryNetwork(addr string, bootNodes []*enode.Node) (*HistoryNetwork, e
 
 	contentQueue := make(chan *discover.ContentElement, 50)
 
-	portalProtocol, err := discover.NewPortalProtocol(conf, string(portalwire.HistoryNetwork), privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
+	portalProtocol, err := discover.NewPortalProtocol(conf, portalwire.History, privKey, conn, localNode, discV5, &storage.MockStorage{Db: make(map[string][]byte)}, contentQueue)
 	if err != nil {
 		return nil, err
 	}