1 files changed, 46 insertions, 14 deletions

diff --git a/btime.c b/btime.c
index bb0857e..5e151be 100644
--- a/btime.c
+++ b/btime.c
@@ -13,6 +13,7 @@
 #include "btime_int.h"
 
 static int same_addr(struct sockaddr_in *a, struct sockaddr_in *b);
+static void gen_cookie(char *s);
 /*
  * btime --
  * 
@@ -27,15 +28,16 @@ unsigned int
 btime(const char *host)
 {
 	int sd;
-	char msg[BTIME_MSGLEN];
+	char cookie[BTIME_MSGLEN];
+	char response[BTIME_MSGLEN];
 	struct sockaddr_in serv_addr;
 	struct sockaddr_in my_addr;
 	struct sockaddr_in resp_addr;
-	int serv_addr_len;
-	int resp_addr_len;
+	socklen_t serv_addr_len;
+	socklen_t resp_addr_len;
 	struct hostent *hostent;
 	in_addr_t inaddr;
-	unsigned int btime;
+	unsigned int btime = 0;
 	ssize_t nbytes;	
 	int retry;
 
@@ -69,10 +71,11 @@ btime(const char *host)
 		return 0;
 	}
 
-	memset(msg, 0, BTIME_MSGLEN);
+	memset(cookie, 0, BTIME_MSGLEN);
+	gen_cookie(cookie);
 	for (retry = 0; retry < MAX_RETRY; retry++) {
 		serv_addr_len = sizeof serv_addr;
-		if ((nbytes = sendto(sd, &msg, BTIME_MSGLEN, MSG_DONTWAIT,
+		if ((nbytes = sendto(sd, &cookie, BTIME_MSGLEN, MSG_DONTWAIT,
 	       	 (struct sockaddr *)&serv_addr, serv_addr_len)) < 0) {
 			if (errno == EAGAIN) {
 				usleep(retry * 100);
@@ -87,18 +90,34 @@ btime(const char *host)
 	
 
 	memset(&resp_addr, 0, sizeof resp_addr);
-	memset(msg, 0, BTIME_MSGLEN);
+	memset(response, 0, BTIME_MSGLEN);
 	for (retry = 0; retry < MAX_RETRY; retry++) {
 		resp_addr_len = sizeof resp_addr;
-		if ((nbytes = recvfrom(sd, &msg, BTIME_MSGLEN, MSG_DONTWAIT, 
+		if ((nbytes = recvfrom(sd, &response, BTIME_MSGLEN, MSG_DONTWAIT, 
 	         (struct sockaddr *)&resp_addr, &resp_addr_len)) < 0) {
 			if (errno == EAGAIN) {
 				usleep(retry * 100);
+				continue;
 			} else {
 				/* Non EAGAIN error... */
 				break;
 			}
+		}
+
+		if (nbytes == 0) {
+			/* Nothing received, try again */
+			continue;
+		} else if (response[0] == 'B' && response[1] == 'T') {
+			/* Check for new style cookie */
+			if (memcmp(cookie, response, COOKIE_LEN) == 0) {
+				btime = strtoul(response+COOKIE_LEN, NULL, 10);
+				break;
+			} else {
+				fprintf(stderr, "Ignoring invalid cookie\n");
+				continue;
+			}
 		} else if (!same_addr(&serv_addr, &resp_addr)) {
+			/* If no cookie, check peer */
 			char *rstr = strdup(inet_ntoa(resp_addr.sin_addr));
 			char *sstr = strdup(inet_ntoa(serv_addr.sin_addr));
 			fprintf(stderr, "Got response from %s instead of %s\n",
@@ -107,16 +126,12 @@ btime(const char *host)
 			free(sstr);
 			continue;
 		} else {
+			/* If peer matches, accept the bare btime */
+			btime = strtoul(response, (char **)NULL, 10);
 			break;
 		}
 	}
 
-	if (nbytes == 0) {
-		btime = 0;
-	} else {
-		btime = strtoul(msg, (char **)NULL, 10);
-	}
-
 	close(sd);
 
 	return btime;
@@ -141,3 +156,20 @@ same_addr(struct sockaddr_in *a, struct sockaddr_in *b)
 		return 1;
 	}
 }
+/* gen_cookie
+ *
+ * 	Generate a random string that btimed will send back to us.
+ */
+static void gen_cookie(char *s)
+{
+	char *a = s;
+	int i;
+
+	*a++ = 'B';
+	*a++ = 'T';
+	for (i = 0; i < COOKIE_RANDOM_PARTS; i++) {
+		*(int *)a = random();
+		a += sizeof(int);
+	}
+	*a = '\n';
+}