Fix: defer initiating GeoIP session cache expiry

- When the tunnel is established, add the GeoIP
  cache entry with no expiry, ensuring it remains
  available when needed and no TTL extension is
  required.
- Once tunnel stops, mark cache entry for expiry.
  This keeps the entry past the end of the tunnel,
  allowing for untunneled final status requests,
  while also ensuring eventual clean up of the
  cache entry.
- Previous fix, fb26f25, did not cover the case
  of clients using SSH API requests followed by an
  untunneled status request as SSH API requests do
  not consult the GeoIP cache and did not extend
  the TTL of the cache entry.

psiphon/server/geoip.go

@@ -175,21 +175,37 @@ func (geoIP *GeoIPService) Lookup(ipAddress string) GeoIPData {
 	return result
 }
 
+// SetSessionCache adds the sessionID/geoIPData pair to the
+// session cache. This value will not expire; the caller must
+// call MarkSessionCacheToExpire to initiate expiry.
+// Calling SetSessionCache for an existing sessionID will
+// replace the previous value and reset any expiry.
 func (geoIP *GeoIPService) SetSessionCache(sessionID string, geoIPData GeoIPData) {
-	geoIP.sessionCache.Set(sessionID, geoIPData, cache.DefaultExpiration)
+	geoIP.sessionCache.Set(sessionID, geoIPData, cache.NoExpiration)
 }
 
-func (geoIP *GeoIPService) GetSessionCache(
-	sessionID string) GeoIPData {
+// MarkSessionCacheToExpire initiates expiry for an existing
+// session cache entry, if the session ID is found in the cache.
+// Concurrency note: SetSessionCache and MarkSessionCacheToExpire
+// should not be called concurrently for a single session ID.
+func (geoIP *GeoIPService) MarkSessionCacheToExpire(sessionID string) {
+	geoIPData, found := geoIP.sessionCache.Get(sessionID)
+	// Note: potential race condition between Get and Set. In practice,
+	// the tunnel server won't clobber a SetSessionCache value by calling
+	// MarkSessionCacheToExpire concurrently.
+	if found {
+		geoIP.sessionCache.Set(sessionID, geoIPData, cache.DefaultExpiration)
+	}
+}
+
+// GetSessionCache returns the cached GeoIPData for the
+// specified session ID; a blank GeoIPData is returned
+// if the session ID is not found in the cache.
+func (geoIP *GeoIPService) GetSessionCache(sessionID string) GeoIPData {
 	geoIPData, found := geoIP.sessionCache.Get(sessionID)
 	if !found {
 		return NewGeoIPData()
 	}
-	// Extend the TTL for this item. If this fails, the cache entry
-	// value remains the same but the TTL is not extended.
-	// Note: sessionCache.Replace can clobber sessionCache.Set changes,
-	// but this is not a concern in practise.
-	_ = geoIP.sessionCache.Replace(sessionID, geoIPData, cache.DefaultExpiration)
 	return geoIPData.(GeoIPData)
 }
 

psiphon/server/tunnelServer.go

@@ -778,10 +778,12 @@ func (sshClient *sshClient) passwordCallback(conn ssh.ConnMetadata, password []b
 	geoIPData := sshClient.geoIPData
 	sshClient.Unlock()
 
-	// Store the GeoIP data associated with the session ID. This makes the GeoIP data
-	// available to the web server for web transport Psiphon API requests. To allow for
-	// post-tunnel final status requests, the lifetime of cached GeoIP records exceeds
-	// the lifetime of the sshClient, and that's why this distinct session cache exists.
+	// Store the GeoIP data associated with the session ID. This makes
+	// the GeoIP data available to the web server for web API requests.
+	// A cache that's distinct from the sshClient record is used to allow
+	// for or post-tunnel final status requests.
+	// If the client is reconnecting with the same session ID, this call
+	// will undo the expiry set by MarkSessionCacheToExpire.
 	sshClient.sshServer.support.GeoIPService.SetSessionCache(sessionID, geoIPData)
 
 	return nil, nil
@@ -862,8 +864,15 @@ func (sshClient *sshClient) stop() {
 	logFields["peak_concurrent_port_forward_count_udp"] = sshClient.udpTrafficState.peakConcurrentPortForwardCount
 	logFields["total_port_forward_count_udp"] = sshClient.udpTrafficState.totalPortForwardCount
 
+	sessionID := sshClient.sessionID
+
 	sshClient.Unlock()
 
+	// Initiate cleanup of the GeoIP session cache. To allow for post-tunnel
+	// final status requests, the lifetime of cached GeoIP records exceeds the
+	// lifetime of the sshClient.
+	sshClient.sshServer.support.GeoIPService.MarkSessionCacheToExpire(sessionID)
+
 	log.LogRawFieldsWithTimestamp(logFields)
 }