add enrollment reply message

nat_traversal.c

@@ -308,6 +308,16 @@ int init(struct peer_info self, struct sockaddr_in punch_server, client* c) {
         return -1;
     }
 
+    // wait for server reply to get own ID
+    uint32_t peer_id = 0;
+    int n = recv(server_sock, &peer_id, sizeof(uint32_t), 0);
+    if (n != 4) {
+        printf("failed to enroll\n");
+        return -1;
+    }
+
+    printf("enroll successfully, ID: %d", ntohl(peer_id));
+
     // wait for message from punch server in another thread
     pthread_t thread_id;
     pthread_create(&thread_id, NULL, server_notify_handler, (void*)&server_sock);

punch_server.go

@@ -64,7 +64,7 @@ func handleConn(c net.Conn) {
 			var peer nat_info
 			err = binary.Read(c, binary.BigEndian, &peer)
 			if err != nil {
-				return
+				continue
 			}
 
 			fmt.Println("peer enrolled, addr: ", string(peer.Ip[:]), peer.Port, peer.Nat_type)
@@ -76,6 +76,10 @@ func handleConn(c net.Conn) {
 			peers_conn[peerID] = c
 			fmt.Println("new peer, id : ", peerID)
 			m.Unlock()
+			err = binary.Write(c, binary.BigEndian, peerID)
+			if err != nil {
+				continue
+			}
 		case GetPeerInfo:
 			var peer_id uint32
 			binary.Read(c, binary.BigEndian, &peer_id)