Adding experiment support for certificate pinning

mikelikespie · mikelikespie · commit 1cba49a173db · 2012-03-25T19:22:44.000-07:00
diff --git a/SocketRocket/SRWebSocket.h b/SocketRocket/SRWebSocket.h
@@ -15,6 +15,7 @@
 //
 
 #import <Foundation/Foundation.h>
+#import <Security/SecCertificate.h>
 
 typedef enum {
     SR_CONNECTING   = 0,
@@ -60,3 +61,17 @@ extern NSString *const SRWebSocketErrorDomain;
 - (void)webSocket:(SRWebSocket *)webSocket didCloseWithCode:(NSInteger)code reason:(NSString *)reason wasClean:(BOOL)wasClean;
 
 @end
+
+
+@interface NSURLRequest (CertificateAdditions)
+
+@property (nonatomic, retain, readonly) NSArray *SR_SSLPinnedCertificates;
+
+@end
+
+
+@interface NSMutableURLRequest (CertificateAdditions)
+
+@property (nonatomic, retain) NSArray *SR_SSLPinnedCertificates;
+
+@end
diff --git a/SocketRocket/SRWebSocket.m b/SocketRocket/SRWebSocket.m
@@ -488,14 +488,24 @@ - (void)_connectToHost:(NSString *)host port:(NSInteger)port;
     _outputStream = CFBridgingRelease(writeStream);
     _inputStream = CFBridgingRelease(readStream);
     
+    
     if (_secure) {
+        NSMutableDictionary *SSLOptions = [[NSMutableDictionary alloc] init];
+             
         [_outputStream setProperty:(__bridge id)kCFStreamSocketSecurityLevelNegotiatedSSL forKey:(__bridge id)kCFStreamPropertySocketSecurityLevel];
+        
+        // If we're using pinned certs, don't validate the certificate chain
+        if ([_urlRequest SR_SSLPinnedCertificates].count) {
+            [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
+        }
+        
         #if DEBUG
+        [SSLOptions setValue:[NSNumber numberWithBool:NO] forKey:(__bridge id)kCFStreamSSLValidatesCertificateChain];
         NSLog(@"SocketRocket: In debug mode.  Allowing connection to any root cert");
-        [_outputStream setProperty:[NSDictionary dictionaryWithObject:[NSNumber numberWithBool:YES] 
-                                                               forKey:(__bridge id)kCFStreamSSLAllowsAnyRoot]
-                            forKey:(__bridge id)kCFStreamPropertySSLSettings];
         #endif
+        
+        [_outputStream setProperty:SSLOptions
+                            forKey:(__bridge id)kCFStreamPropertySSLSettings];
     }
     
     _inputStream.delegate = self;
@@ -1232,6 +1242,10 @@ - (void)_sendFrameWithOpcode:(SROpCode)opcode data:(id)data;
 - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
 {
     //    SRFastLog(@"%@ Got stream event %d", aStream, eventCode);
+    if (eventCode == NSStreamEventErrorOccurred || eventCode == NSStreamEventErrorOccurred) {
+        
+        
+    }
     dispatch_async(_workQueue, ^{
         switch (eventCode) {
             case NSStreamEventOpenCompleted: {
@@ -1240,6 +1254,35 @@ - (void)stream:(NSStream *)aStream handleEvent:(NSStreamEvent)eventCode;
                     return;
                 }
                 
+                
+                if (_secure && aStream == _outputStream) {
+                    NSArray *sslCerts = [_urlRequest SR_SSLPinnedCertificates];
+                    if (sslCerts) {
+                        SecTrustRef secTrust = (__bridge SecTrustRef)[aStream propertyForKey:(__bridge id)kCFStreamPropertySSLPeerTrust];
+                        BOOL found = NO;
+                        if (secTrust) {
+                            NSInteger numCerts = SecTrustGetCertificateCount(secTrust);
+                            for (NSInteger i = 0; i < numCerts && !found; i++) {
+                                SecCertificateRef cert = SecTrustGetCertificateAtIndex(secTrust, i);
+                                NSData *certData = CFBridgingRelease(SecCertificateCopyData(cert));
+                                
+                                for (id ref in sslCerts) {
+                                    SecCertificateRef trustedCert = (__bridge SecCertificateRef)ref;
+                                    NSData *trustedCertData = CFBridgingRelease(SecCertificateCopyData(trustedCert));
+                                    
+                                    if ([trustedCertData isEqualToData:certData]) {
+                                        found = YES;
+                                    }
+                                }
+                            }
+                        }
+                        
+                        if (!found) {
+                            [self _failWithError:[NSError errorWithDomain:@"org.lolrus.SocketRocket" code:23556 userInfo:[NSDictionary dictionaryWithObject:[NSString stringWithFormat:@"Invalid server cert"] forKey:NSLocalizedDescriptionKey]]];
+                        }
+                    }
+                    
+                }
                 assert(_readBuffer);
                 
                 if (self.readyState == SR_CONNECTING && aStream == _inputStream) {
@@ -1347,4 +1390,26 @@ - (id)initWithScanner:(stream_scanner)scanner handler:(data_callback)handler byt
 @end
 
 
+@implementation  NSURLRequest (CertificateAdditions)
+
+- (NSArray *)SR_SSLPinnedCertificates;
+{
+    return [NSURLProtocol propertyForKey:@"SR_SSLPinnedCertificates" inRequest:self];
+}
+
+@end
+
+@implementation  NSMutableURLRequest (CertificateAdditions)
+
+- (NSArray *)SR_SSLPinnedCertificates;
+{
+    return [NSURLProtocol propertyForKey:@"SR_SSLPinnedCertificates" inRequest:self];
+}
+
+- (void)setSR_SSLPinnedCertificates:(NSArray *)SR_SSLPinnedCertificates;
+{
+    [NSURLProtocol setProperty:SR_SSLPinnedCertificates forKey:@"SR_SSLPinnedCertificates" inRequest:self];
+}
+
+@end
 
