From 1f37980a0dd3980595bc1231d64772f18c0265ca Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Thu, 11 Oct 2018 10:44:28 -0400
Subject: [PATCH] Suppress UD against production service to avoid
 de-registration.

---
 .../src/Messages/UD/OWSUDManager.swift             | 10 ++++++++++
 .../src/Network/WebSockets/OWSWebSocket.m          |  8 +++++---
 SignalServiceKit/src/TSConstants.h                 |  5 +++--
 SignalServiceKit/src/TSConstants.m                 | 14 ++++++++++++++
 SignalServiceKit/src/TSPrefix.h                    |  3 ++-
 5 files changed, 34 insertions(+), 6 deletions(-)
 create mode 100644 SignalServiceKit/src/TSConstants.m

diff --git a/SignalServiceKit/src/Messages/UD/OWSUDManager.swift b/SignalServiceKit/src/Messages/UD/OWSUDManager.swift
index d2374ea2f..825a1c7f3 100644
--- a/SignalServiceKit/src/Messages/UD/OWSUDManager.swift
+++ b/SignalServiceKit/src/Messages/UD/OWSUDManager.swift
@@ -26,6 +26,8 @@ public enum UnidentifiedAccessMode: Int {
 
     @objc func trustRoot() -> ECPublicKey
 
+    @objc func isUDEnabled() -> Bool
+
     // MARK: - Recipient State
 
     @objc
@@ -173,6 +175,9 @@ public class OWSUDManagerImpl: NSObject, OWSUDManager {
     // if we have a valid profile key for them.
     @objc
     public func udAccessKeyForRecipient(_ recipientId: RecipientIdentifier) -> SMKUDAccessKey? {
+        guard isUDEnabled() else {
+            return nil
+        }
         let theirAccessMode = unidentifiedAccessMode(recipientId: recipientId)
         if theirAccessMode == .unrestricted {
             return SMKUDAccessKey(randomKeyData: ())
@@ -281,6 +286,11 @@ public class OWSUDManagerImpl: NSObject, OWSUDManager {
         }
     }
 
+    @objc
+    public func isUDEnabled() -> Bool {
+        return !IsUsingProductionService()
+    }
+
     @objc
     public func trustRoot() -> ECPublicKey {
         guard let trustRootData = NSData(fromBase64String: kUDTrustRoot) else {
diff --git a/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.m b/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.m
index 968748a0d..1d09c291b 100644
--- a/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.m
+++ b/SignalServiceKit/src/Network/WebSockets/OWSWebSocket.m
@@ -264,6 +264,10 @@ NSString *const kNSNotification_OWSWebSocketStateDidChange = @"kNSNotification_O
     return OWSWebsocketSecurityPolicy.sharedPolicy;
 }
 
+- (id<OWSUDManager>)udManager {
+    return SSKEnvironment.shared.udManager;
+}
+
 #pragma mark -
 
 // We want to observe these notifications lazily to avoid accessing
@@ -1024,12 +1028,10 @@ NSString *const kNSNotification_OWSWebSocketStateDidChange = @"kNSNotification_O
     }
 #endif
 
-#ifndef UD_ENABLED
-    if (self.webSocketType == OWSWebSocketTypeUD) {
+    if (!self.udManager.isUDEnabled && self.webSocketType == OWSWebSocketTypeUD) {
         OWSLogWarn(@"Suppressing UD socket in prod.");
         return;
     }
-#endif
 
     if (!AppReadiness.isAppReady) {
         static dispatch_once_t onceToken;
diff --git a/SignalServiceKit/src/TSConstants.h b/SignalServiceKit/src/TSConstants.h
index 9fb966355..ed72036f7 100644
--- a/SignalServiceKit/src/TSConstants.h
+++ b/SignalServiceKit/src/TSConstants.h
@@ -37,7 +37,7 @@ typedef NS_ENUM(NSInteger, TSWhisperMessageType) {
 //#define contactDiscoveryURL @"https://api.directory.signal.org"
 //// TODO: The production value is not yet known.
 //#define kUDTrustRoot @"BbqY1DzohE4NUZoVF+L18oUPrK3kILllLEJh2UnPSsEx"
-////#define UD_ENABLED
+//#define USING_PRODUCTION_SERVICE
 
 //#else
 
@@ -49,10 +49,11 @@ typedef NS_ENUM(NSInteger, TSWhisperMessageType) {
 #define textSecureCDNReflectorHost @"meek-signal-cdn-staging.appspot.com";
 #define contactDiscoveryURL @"https://api-staging.directory.signal.org"
 #define kUDTrustRoot @"BbqY1DzohE4NUZoVF+L18oUPrK3kILllLEJh2UnPSsEx"
-#define UD_ENABLED
 
 //#endif
 
+BOOL IsUsingProductionService(void);
+
 #define textSecureAccountsAPI @"v1/accounts"
 #define textSecureAttributesAPI @"/attributes/"
 
diff --git a/SignalServiceKit/src/TSConstants.m b/SignalServiceKit/src/TSConstants.m
new file mode 100644
index 000000000..b5b92adef
--- /dev/null
+++ b/SignalServiceKit/src/TSConstants.m
@@ -0,0 +1,14 @@
+//
+//  Copyright (c) 2018 Open Whisper Systems. All rights reserved.
+//
+
+#import "TSConstants.h"
+
+BOOL IsUsingProductionService()
+{
+#ifdef USING_PRODUCTION_SERVICE
+    return YES;
+#else
+    return NO;
+#endif
+}
diff --git a/SignalServiceKit/src/TSPrefix.h b/SignalServiceKit/src/TSPrefix.h
index c9c36c4e5..d482637cb 100644
--- a/SignalServiceKit/src/TSPrefix.h
+++ b/SignalServiceKit/src/TSPrefix.h
@@ -12,6 +12,7 @@ static const NSUInteger ddLogLevel = DDLogLevelAll;
 static const NSUInteger ddLogLevel = DDLogLevelInfo;
 #endif
 #import "OWSAnalytics.h"
+#import "SSKAsserts.h"
+#import "TSConstants.h"
 #import <SignalCoreKit/NSObject+OWS.h>
 #import <SignalCoreKit/OWSAsserts.h>
-#import "SSKAsserts.h"