gui: serve country code

Author: mmcgee-jump

book/api/websocket.md

@@ -1400,7 +1400,8 @@ identity is no longer in these three data sources, it will be removed.
                     "gossip": "93.119.195.160:8001",
                     "tpu": "192.64.85.26:8000",
                     // ... other sockets ...
-                }
+                },
+                "country_code": "CN"
             },
             "vote": [
                 {
@@ -1438,6 +1439,7 @@ identity is no longer in these three data sources, it will be removed.
 | version       | `string\|null` | Software version being advertised by the validator. Might be `null` if the validator is not gossiping a version, or we have received the contact information but not the version yet. The version string, if not null, will always be formatted like `major`.`minor`.`patch` where `major`, `minor`, and `patch` are `u16`s |
 | feature_set   | `number\|null` | First four bytes of the `FeatureSet` hash interpreted as a little endian `u32`. Might be `null` if the validator is not gossiping a feature set, or we have received the contact information but not the feature set yet |
 | sockets       | `[key: string]: string` | A dictionary of sockets that are advertised by the validator. `key` will be one of gossip `serve_repair_quic`, `rpc`, `rpc_pubsub`, `serve_repair`, `tpu`, `tpu_forwards`, `tpu_forwards_quic`, `tpu_quic`, `tpu_vote`, `tvu`, `tvu_quic`, `tpu_vote_quic`, or `alpenglow`. The value is an address like `<addr>:<port>`: the location to send traffic to for this validator with the given protocol. Address might be either an IPv4 or an IPv6 address |
+| country_code  | `string\|null` | ISO 3166-1 alpha-2 country code of where the validator is located, determined by GeoIP lookup on the gossip IP address. Country code may not be correct and is a best estimate. If no country code could be determined, will be `null`. |
 
 **`PeerUpdateVoteAccount`**
 | Field           | Type           | Description |

src/disco/gui/Local.mk

@@ -41,3 +41,4 @@ FD_GUI_FRONTEND_DEV_GZ_FILES := $(patsubst src/disco/gui/dist_dev/%, src/disco/g
 FD_GUI_FRONTEND_DEV_ZST_FILES := $(patsubst src/disco/gui/dist_dev/%, src/disco/gui/dist_dev_cmp/%.zst, $(FD_GUI_FRONTEND_DEV_FILES))
 
 $(OBJDIR)/obj/disco/gui/generated/http_import_dist.d: $(FD_GUI_FRONTEND_STABLE_GZ_FILES) $(FD_GUI_FRONTEND_STABLE_ZST_FILES) $(FD_GUI_FRONTEND_ALPHA_GZ_FILES) $(FD_GUI_FRONTEND_ALPHA_ZST_FILES) $(FD_GUI_FRONTEND_DEV_GZ_FILES) $(FD_GUI_FRONTEND_DEV_ZST_FILES)
+$(OBJDIR)/obj/disco/gui/fd_gui.d: src/disco/gui/ipinfo.bin.zstd

src/disco/gui/convert_ipinfo.py

@@ -0,0 +1,55 @@
+#!/usr/bin/env python3
+import struct
+import ipaddress
+import csv
+import argparse
+import zstandard as zstd
+
+def main():
+    parser = argparse.ArgumentParser(description='Convert IPInfo CSV to binary format')
+    parser.add_argument('input', help='Input CSV file path')
+    parser.add_argument('output', help='Output binary file path')
+    args = parser.parse_args()
+
+    country_codes = set()
+    with open(args.input, 'r') as r:
+        reader = csv.DictReader(r)
+        for row in reader:
+            try:
+                ipaddress.IPv4Network(row['network'])
+            except ipaddress.AddressValueError:
+                continue
+            assert len(row['country_code']) == 2
+            country_codes.add(row['country_code'])
+
+    assert len(country_codes) < 256, f"Too many country codes ({len(country_codes)}) to fit in a byte (max 255)"
+
+    country_to_index = {cc: idx for idx, cc in enumerate(sorted(country_codes))}
+
+    with open(args.input, 'r') as r:
+        reader = csv.DictReader(r)
+        with open(args.output, 'wb') as f:
+            cctx = zstd.ZstdCompressor(level=22)
+            with cctx.stream_writer(f) as w:
+                w.write(struct.pack('<Q', len(country_codes)))
+                for cc in sorted(country_codes):
+                    w.write(cc.encode('ascii'))
+
+                records = 0
+                for row in reader:
+                    try:
+                        network = ipaddress.IPv4Network(row['network'])
+                    except ipaddress.AddressValueError:
+                        continue
+
+                    country_idx = country_to_index[row['country_code']]
+
+                    w.write(struct.pack('>I', int(network.network_address)))
+                    w.write(struct.pack('<B', network.prefixlen))
+                    w.write(struct.pack('<B', country_idx))
+                    records += 1
+
+    print(f"Converted {records} records with {len(country_codes)} country codes to {args.output}")
+
+if __name__ == "__main__":
+    main()

src/disco/gui/dump.py

@@ -0,0 +1,11 @@
+import asyncio
+import websockets
+import json
+
+async def dump_messages(uri):
+    async with websockets.connect(uri, max_size=1_000_000_000) as websocket:
+        while True:
+            frame = await websocket.recv()
+            print(json.dumps(json.loads(frame)))
+
+asyncio.get_event_loop().run_until_complete(dump_messages('ws://localhost:80/websocket'))

src/disco/gui/fd_gui.c

@@ -1,4 +1,5 @@
 #include "fd_gui.h"
+#include "fd_gui_peers.h"
 #include "fd_gui_printf.h"
 
 #include "../metrics/fd_metrics.h"
@@ -10,18 +11,111 @@
 #include "../../disco/pack/fd_pack.h"
 #include "../../disco/pack/fd_pack_cost.h"
 
+FD_IMPORT_BINARY( ipinfo, "src/disco/gui/ipinfo.bin.zstd" );
+
 #include <stdio.h>
 
+#if FD_HAS_ZSTD
+#define FD_HTTP_ZSTD_COMPRESSION_LEVEL 3
+#define ZSTD_STATIC_LINKING_ONLY
+#include <zstd.h>
+/* TODO: Just use a small buffer with streaming decompression */
+#define IPINFO_DECOMPRESSED_SZ (1UL<<24UL) /* 16 MiB */
+#define IPINFO_MAX_NODES       (1UL<<22UL) /* 4M nodes */
+#endif
+
 FD_FN_CONST ulong
 fd_gui_align( void ) {
   return 128UL;
 }
 
 FD_FN_CONST ulong
 fd_gui_footprint( void ) {
-  return sizeof(fd_gui_t);
+  ulong l = FD_LAYOUT_INIT;
+  l = FD_LAYOUT_APPEND( l, fd_gui_align(),                sizeof( fd_gui_t ) );
+#if FD_HAS_ZSTD
+  l = FD_LAYOUT_APPEND( l, 1UL,                           IPINFO_DECOMPRESSED_SZ );
+  l = FD_LAYOUT_APPEND( l, alignof(fd_gui_ipinfo_node_t), sizeof(fd_gui_ipinfo_node_t)*IPINFO_MAX_NODES );
+#endif
+  return FD_LAYOUT_FINI( l, fd_gui_align() );
 }
 
+#if FD_HAS_ZSTD
+
+static void
+build_ipinfo_trie( fd_gui_t *             gui,
+                   uchar *                ipinfo_buffer,
+                   fd_gui_ipinfo_node_t * nodes ) {
+  gui->ipinfo.nodes = nodes;
+
+  ulong actual_sz = ZSTD_decompress( ipinfo_buffer, IPINFO_DECOMPRESSED_SZ, ipinfo, ipinfo_sz );
+  FD_TEST( !ZSTD_isError( actual_sz ) );
+  FD_TEST( actual_sz>8UL );
+
+  ulong country_code_cnt = FD_LOAD( ulong, ipinfo_buffer );
+  FD_TEST( country_code_cnt && country_code_cnt<256UL ); /* 256 reserved for unknown */
+  FD_TEST( actual_sz>=8UL+country_code_cnt*2UL );
+
+  for( ulong i=0UL; i<country_code_cnt; i++ ) {
+    fd_memcpy( gui->ipinfo.country_code[ i ], ipinfo_buffer+8UL+i*2UL, 2UL );
+    gui->ipinfo.country_code[ i ][ 2 ] = '\0';
+  }
+
+  ulong processed = 8UL+country_code_cnt*2UL;
+  FD_TEST( !((actual_sz-processed)%6UL) );
+  FD_TEST( (actual_sz-processed)/6UL<=IPINFO_MAX_NODES-1UL );
+
+  fd_gui_ipinfo_node_t * root = &nodes[ 0 ];
+  root->left = NULL;
+  root->right = NULL;
+  root->has_prefix = 0;
+
+  ulong node_cnt = 1UL;
+  while( processed<actual_sz ) {
+    uint ip_addr = fd_uint_bswap( FD_LOAD( uint, ipinfo_buffer+processed ) );
+    uchar prefix_len = *( ipinfo_buffer+processed+4UL );
+    FD_TEST( prefix_len<=32UL );
+    uchar country_idx = *( ipinfo_buffer+processed+5UL );
+    FD_TEST( country_idx<country_code_cnt );
+
+    fd_gui_ipinfo_node_t * node = root;
+    for( uchar bit_pos=0; bit_pos<prefix_len; bit_pos++ ) {
+      uchar bit = (ip_addr >> (31 - bit_pos)) & 1;
+
+      fd_gui_ipinfo_node_t * child;
+      if( FD_LIKELY( !bit ) ) {
+        child = node->left;
+        if( FD_LIKELY( !child ) ) {
+          FD_TEST( node_cnt<IPINFO_MAX_NODES );
+          child = &nodes[ node_cnt++ ];
+          child->left = NULL;
+          child->right = NULL;
+          child->has_prefix = 0;
+          node->left = child;
+        }
+      } else {
+        child = node->right;
+        if( FD_LIKELY( !child ) ) {
+          FD_TEST( node_cnt<IPINFO_MAX_NODES );
+          child = &nodes[ node_cnt++ ];
+          child->left = NULL;
+          child->right = NULL;
+          child->has_prefix = 0;
+          node->right = child;
+        }
+      }
+      node = child;
+    }
+
+    node->has_prefix = 1;
+    node->country_code_idx = country_idx;
+
+    processed += 6UL;
+  }
+}
+
+#endif
+
 void *
 fd_gui_new( void *                shmem,
             fd_http_server_t *    http,
@@ -52,7 +146,12 @@ fd_gui_new( void *                shmem,
     return NULL;
   }
 
-  fd_gui_t * gui = (fd_gui_t *)shmem;
+  FD_SCRATCH_ALLOC_INIT( l, shmem );
+  fd_gui_t * gui                = FD_SCRATCH_ALLOC_APPEND( l, fd_gui_align(),                sizeof(fd_gui_t) );
+#if FD_HAS_ZSTD
+  uchar * _ipinfo_buffer        = FD_SCRATCH_ALLOC_APPEND( l, 1UL,                           IPINFO_DECOMPRESSED_SZ );
+  fd_gui_ipinfo_node_t * _nodes = FD_SCRATCH_ALLOC_APPEND( l, alignof(fd_gui_ipinfo_node_t), sizeof(fd_gui_ipinfo_node_t)*IPINFO_MAX_NODES );
+#endif
 
   gui->http = http;
   gui->topo = topo;
@@ -154,7 +253,6 @@ fd_gui_new( void *                shmem,
   for( ulong i=0UL; i<FD_GUI_LEADER_CNT; i++ ) gui->leader_slots[ i ]->slot      = ULONG_MAX;
   gui->leader_slots_cnt      = 0UL;
 
-
   gui->block_engine.has_block_engine = 0;
 
   gui->epoch.has_epoch[ 0 ] = 0;
@@ -174,6 +272,10 @@ fd_gui_new( void *                shmem,
   gui->summary.catch_up_repair_sz  = 0UL;
   gui->summary.catch_up_turbine_sz = 0UL;
 
+#if FD_HAS_ZSTD
+  build_ipinfo_trie( gui, _ipinfo_buffer, _nodes );
+#endif
+
   return gui;
 }
 

src/disco/gui/fd_gui.h

@@ -372,8 +372,8 @@ struct fd_gui_slot_rankings {
 typedef struct fd_gui_slot_rankings fd_gui_slot_rankings_t;
 
 struct fd_gui_ephemeral_slot {
-      ulong slot; /* ULONG_MAX indicates invalid/evicted */
-      long timestamp_arrival_nanos;
+  ulong slot; /* ULONG_MAX indicates invalid/evicted */
+  long timestamp_arrival_nanos;
 };
 typedef struct fd_gui_ephemeral_slot fd_gui_ephemeral_slot_t;
 
@@ -650,6 +650,11 @@ struct fd_gui {
     fd_gui_scheduler_counts_t scheduler_counts_snap[ FD_GUI_SCHEDULER_COUNT_SNAP_CNT ][ 1 ];
   } summary;
 
+  struct {
+    fd_gui_ipinfo_node_t * nodes;
+    char country_code[ 512 ][ 3 ]; /* ISO 3166-1 alpha-2 country codes */
+  } ipinfo;
+
   fd_gui_slot_t slots[ FD_GUI_SLOTS_CNT ][ 1 ];
 
   fd_gui_leader_slot_t leader_slots[ FD_GUI_LEADER_CNT ][ 1 ];