/*
 * Endless
 * Copyright (c) 2014-2015 joshua stein <jcs@jcs.org>
 *
 * Originally created by Mike Abdullah on 17/03/2009.
 * Copyright 2009 Karelia Software. All rights reserved.
 *
 * Originally from ConnectionKit 2.0 branch; source at:
 * http://www.opensource.utr-software.com/source/connection/branches/2.0/CKHTTPConnection.m
 * (CKHTTPConnection.m last updated rev 1242, 2009-06-16 09:40:21 -0700, by mabdullah)
 *
 * Under Modified BSD License, as per description at
 * http://www.opensource.utr-software.com/
 */

#import "CKHTTPConnection.h"
#import "HostSettings.h"
#import "SSLCertificate.h"

@interface CKHTTPConnection ()
- (CFHTTPMessageRef)HTTPRequest;
- (NSInputStream *)HTTPStream;
- (void)start;
- (id <CKHTTPConnectionDelegate>)delegate;
@end

@interface CKHTTPConnection (Authentication) <NSURLAuthenticationChallengeSender>
- (CKHTTPAuthenticationChallenge *)currentAuthenticationChallenge;
@end

@interface CKHTTPAuthenticationChallenge : NSURLAuthenticationChallenge
{
	CFHTTPAuthenticationRef _HTTPAuthentication;
}

- (id)initWithResponse:(CFHTTPMessageRef)response proposedCredential:(NSURLCredential *)credential previousFailureCount:(NSInteger)failureCount failureResponse:(NSHTTPURLResponse *)URLResponse sender:(id <NSURLAuthenticationChallengeSender>)sender;
- (CFHTTPAuthenticationRef)CFHTTPAuthentication;

@end


@implementation CKHTTPConnection

+ (CKHTTPConnection *)connectionWithRequest:(NSURLRequest *)request delegate:(id <CKHTTPConnectionDelegate>)delegate
{
	return [[self alloc] initWithRequest:request delegate:delegate];
}

- (id)initWithRequest:(NSURLRequest *)request delegate:(id <CKHTTPConnectionDelegate>)delegate;
{
	NSParameterAssert(request);

	if (!(self = [super init]))
		return nil;
	
	_delegate = delegate;
	_HTTPRequest = [request makeHTTPMessage];
	_HTTPBodyStream = [request HTTPBodyStream];

	[self start];
	
	return self;
}

- (void)dealloc
{
	CFRelease(_HTTPRequest);
}

#pragma mark Accessors

- (CFHTTPMessageRef)HTTPRequest {
	return _HTTPRequest;
}

- (NSInputStream *)HTTPStream {
	return _HTTPStream;
}

- (NSInputStream *)stream {
	return (NSInputStream *)[self HTTPStream];
}

- (id <CKHTTPConnectionDelegate>)delegate {
	return _delegate;
}

#pragma mark Status handling

- (void)start
{
	NSAssert(!_HTTPStream, @"Connection already started");
	HostSettings *hs;
	
	if (_HTTPBodyStream)
		_HTTPStream = (__bridge_transfer NSInputStream *)(CFReadStreamCreateForStreamedHTTPRequest(NULL, [self HTTPRequest], (__bridge CFReadStreamRef)_HTTPBodyStream));
	else
		_HTTPStream = (__bridge_transfer NSInputStream *)CFReadStreamCreateForHTTPRequest(NULL, [self HTTPRequest]);
	
	/* we're handling redirects ourselves */
	CFReadStreamSetProperty((__bridge CFReadStreamRef)(_HTTPStream), kCFStreamPropertyHTTPShouldAutoredirect, kCFBooleanFalse);

	NSString *method = (__bridge_transfer NSString *)CFHTTPMessageCopyRequestMethod([self HTTPRequest]);
	if ([[method uppercaseString] isEqualToString:@"GET"])
		CFReadStreamSetProperty((__bridge CFReadStreamRef)(_HTTPStream), kCFStreamPropertyHTTPAttemptPersistentConnection, kCFBooleanTrue);
	else
		CFReadStreamSetProperty((__bridge CFReadStreamRef)(_HTTPStream), kCFStreamPropertyHTTPAttemptPersistentConnection, kCFBooleanFalse);
	
	/* set SSL protocol version enforcement before opening, when using kCFStreamSSLLevel */
	NSURL *url = (__bridge_transfer NSURL *)(CFHTTPMessageCopyRequestURL([self HTTPRequest]));
	if ([[[url scheme] lowercaseString] isEqualToString:@"https"]) {
		hs = [HostSettings settingsOrDefaultsForHost:[url host]];

		CFMutableDictionaryRef sslOptions = CFDictionaryCreateMutable(kCFAllocatorDefault, 0, &kCFTypeDictionaryKeyCallBacks, &kCFTypeDictionaryValueCallBacks);

		BOOL setOptions = NO;

		if ([hs boolSettingOrDefault:HOST_SETTINGS_KEY_IGNORE_TLS_ERRORS]
			&& [[NSUserDefaults standardUserDefaults] boolForKey:@"allow_tls_error_ignore"])
		{
			CFDictionarySetValue(sslOptions, kCFStreamSSLValidatesCertificateChain, kCFBooleanFalse);
			setOptions = YES;
		}
		
		if ([[hs settingOrDefault:HOST_SETTINGS_KEY_TLS] isEqualToString:HOST_SETTINGS_TLS_12]) {
			/* kTLSProtocol12 allows lower protocols, so use kCFStreamSSLLevel to force 1.2 */
			
			CFDictionarySetValue(sslOptions, kCFStreamSSLLevel, CFSTR("kCFStreamSocketSecurityLevelTLSv1_2"));
			setOptions = YES;

#ifdef TRACE_HOST_SETTINGS
			NSLog(@"[HostSettings] set TLS/SSL min level for %@ to TLS 1.2", [url host]);
#endif
		}

		if (setOptions) {
			CFReadStreamSetProperty((__bridge CFReadStreamRef)_HTTPStream, kCFStreamPropertySSLSettings, sslOptions);
		}
	}
	
	[_HTTPStream setDelegate:(id<NSStreamDelegate>)self];
	[_HTTPStream scheduleInRunLoop:[NSRunLoop currentRunLoop] forMode:NSDefaultRunLoopMode];
	[_HTTPStream open];

#ifdef TRACE_RAW_HTTP
	NSData *d = (__bridge NSData *)CFHTTPMessageCopySerializedMessage((__bridge CFHTTPMessageRef _Nonnull)([_HTTPStream propertyForKey:(NSString *)kCFStreamPropertyHTTPFinalRequest]));
	NSLog(@"[CKHTTPConnection] final request for %@\n%@", url, [[[NSString alloc] initWithBytes:[d bytes] length:[d length] encoding:NSUTF8StringEncoding] stringByReplacingOccurrencesOfString:@"\r" withString:@""]);
#endif
	
	/* for other SSL options, these need an SSLContextRef which doesn't exist until the stream is opened */
	if ([[[url scheme] lowercaseString] isEqualToString:@"https"]) {
		SSLContextRef sslContext = (__bridge SSLContextRef)[_HTTPStream propertyForKey:(__bridge NSString *)kCFStreamPropertySSLContext];
		if (sslContext != NULL) {
			SSLSessionState sslState;
			SSLGetSessionState(sslContext, &sslState);
			
			/* if we're not idle, this is probably a persistent connection we already opened and negotiated */
			if (sslState == kSSLIdle) {
				if (![self disableWeakSSLCiphers:sslContext]) {
					NSLog(@"[CKHTTPConnection] failed disabling weak ciphers, aborting connection");
					[self _cancelStream];
					return;
				}
			}
		}
	}
}

- (BOOL)disableWeakSSLCiphers:(SSLContextRef)sslContext
{
	OSStatus status;
	size_t numSupported;
	SSLCipherSuite *supported = NULL;
	SSLCipherSuite *enabled = NULL;
	int numEnabled = 0;
	
	status = SSLGetNumberSupportedCiphers(sslContext, &numSupported);
	if (status != noErr) {
		NSLog(@"[CKHTTPConnection] failed getting number of supported ciphers");
		return NO;
	}
	
	supported = (SSLCipherSuite *)malloc(numSupported * sizeof(SSLCipherSuite));
	status = SSLGetSupportedCiphers(sslContext, supported, &numSupported);
	if (status != noErr) {
		NSLog(@"[CKHTTPConnection] failed getting supported ciphers");
		free(supported);
		return NO;
	}
	
	enabled = (SSLCipherSuite *)malloc(numSupported * sizeof(SSLCipherSuite));
	
	for (int i = 0; i < numSupported; i++) {
		switch (supported[i]) {
		case TLS_ECDHE_RSA_WITH_CHACHA20_POLY1305_SHA256:
		case TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256:
		case TLS_ECDHE_RSA_WITH_AES_256_GCM_SHA384:
		case TLS_DHE_RSA_WITH_AES_256_GCM_SHA384:
		case TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA384:
		case TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256:
		case TLS_ECDHE_RSA_WITH_AES_128_GCM_SHA256:
		case TLS_ECDHE_ECDSA_WITH_AES_256_CBC_SHA:
		case TLS_ECDHE_ECDSA_WITH_AES_128_CBC_SHA:
		case TLS_ECDHE_RSA_WITH_AES_128_CBC_SHA:
		case TLS_ECDHE_RSA_WITH_AES_256_CBC_SHA:
		case TLS_RSA_WITH_AES_128_CBC_SHA:
		case TLS_RSA_WITH_AES_256_CBC_SHA:
			enabled[numEnabled++] = supported[i];
			break;
		default:
			if (supported[i] > TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305_SHA256 &&
			    supported[i] < TLS_EMPTY_RENEGOTIATION_INFO_SCSV) {
				/* assume anything newer than what we've allowed is more secure */
				enabled[numEnabled++] = supported[i];
				break;
			}
#ifdef TRACE
			NSLog(@"[CKHTTPConnection] disabling weak SSL/TLS cipher 0x%d", supported[i]);
#endif
		}
	}
	free(supported);

	status = SSLSetEnabledCiphers(sslContext, enabled, numEnabled);
	free(enabled);
	if (status != noErr) {
		NSLog(@"[CKHTTPConnection] failed setting enabled ciphers on %@: %d", sslContext, (int)status);
		return NO;
	}
	
	return YES;
}

- (void)_cancelStream
{
	// Support method to cancel the HTTP stream, but not change the delegate. Used for:
	//  A) Cancelling the connection
	//  B) Waiting to restart the connection while authentication takes place
	//  C) Restarting the connection after an HTTP redirect
	[_HTTPStream close];
	CFBridgingRelease((__bridge_retained CFTypeRef)(_HTTPStream));
	//[_HTTPStream release]; 
	_HTTPStream = nil;
}

- (void)cancel
{
	// Cancel the stream and stop the delegate receiving any more info
	[self _cancelStream];
	_delegate = nil;
}

- (void)stream:(NSInputStream *)theStream handleEvent:(NSStreamEvent)streamEvent
{
	NSParameterAssert(theStream == [self stream]);
	
	NSURL *URL = [theStream propertyForKey:(NSString *)kCFStreamPropertyHTTPFinalURL];

	if (!_haveReceivedResponse) {
		CFHTTPMessageRef response = (__bridge CFHTTPMessageRef)[theStream propertyForKey:(NSString *)kCFStreamPropertyHTTPResponseHeader];
		if (response && CFHTTPMessageIsHeaderComplete(response)) {
			NSHTTPURLResponse *URLResponse = [[NSHTTPURLResponse alloc] initWithURL:URL statusCode:CFHTTPMessageGetResponseStatusCode(response) HTTPVersion:(__bridge NSString * _Nullable)(CFHTTPMessageCopyVersion(response)) headerFields:(__bridge NSDictionary<NSString *,NSString *> * _Nullable)(CFHTTPMessageCopyAllHeaderFields(response))];
			
			NSData *d = (__bridge NSData *)CFHTTPMessageCopySerializedMessage((__bridge CFHTTPMessageRef _Nonnull)([_HTTPStream propertyForKey:(NSString *)kCFStreamPropertyHTTPResponseHeader]));
			
			/* work around bug where CFHTTPMessageIsHeaderComplete reports true but there is no actual header data to be found */
			if ([d length] < 5) {
#ifdef TRACE
				NSLog(@"[CKHTTPConnection] hit CFHTTPMessageIsHeaderComplete bug, waiting for more data");
#endif
				goto process;
			}
			
#ifdef TRACE_RAW_HTTP
			NSLog(@"[CKHTTPConnection] final response for %@\n%@", URL, [[[NSString alloc] initWithBytes:[d bytes] length:[d length] encoding:NSUTF8StringEncoding] stringByReplacingOccurrencesOfString:@"\r" withString:@""]);
#endif
			
			// If the response was an authentication failure, try to request fresh credentials.
			if ([URLResponse statusCode] == 401 || [URLResponse statusCode] == 407) {
				// Cancel any further loading and ask the delegate for authentication
				[self _cancelStream];
				
				NSAssert(![self currentAuthenticationChallenge], @"Authentication challenge received while another is in progress");
				
				_authenticationChallenge = [[CKHTTPAuthenticationChallenge alloc] initWithResponse:response proposedCredential:nil previousFailureCount:_authenticationAttempts failureResponse:URLResponse sender:self];
				
				if ([self currentAuthenticationChallenge]) {
					_authenticationAttempts++;
					[[self delegate] HTTPConnection:self didReceiveAuthenticationChallenge:[self currentAuthenticationChallenge]];
					return; // Stops the delegate being sent a response received message
				}
			}
			
			// By reaching this point, the response was not a valid request for authentication,
			// so go ahead and report it
			_haveReceivedResponse = YES;
			[[self delegate] HTTPConnection:self didReceiveResponse:URLResponse];
		}
	}
	
process:
	switch (streamEvent) {
	case NSStreamEventOpenCompleted:
		break;
	case NSStreamEventHasSpaceAvailable:
		socketReady = true;
		break;
	case NSStreamEventErrorOccurred:
		if (!socketReady && !retriedSocket) {
			/* probably a dead keep-alive socket from the get go */
			retriedSocket = true;
#ifdef TRACE
			NSLog(@"[CKHTTPConnection] socket for %@ dead but never writable, retrying (%@)", [URL absoluteString], [theStream streamError]);
#endif
			[self _cancelStream];
			[self start];
		}
		else
			[[self delegate] HTTPConnection:self didFailWithError:[theStream streamError]];
		break;

	case NSStreamEventEndEncountered:   // Report the end of the stream to the delegate
		[[self delegate] HTTPConnectionDidFinishLoading:self];
		break;

	case NSStreamEventHasBytesAvailable: {
		socketReady = true;
		
		if ([[[URL scheme] lowercaseString] isEqualToString:@"https"]) {
			SecTrustRef trust = (__bridge SecTrustRef)[theStream propertyForKey:(__bridge NSString *)kCFStreamPropertySSLPeerTrust];
			if (trust != nil) {
				SSLCertificate *cert = [[SSLCertificate alloc] initWithSecTrustRef:trust];
				
				SSLContextRef sslContext = (__bridge SSLContextRef)[theStream propertyForKey:(__bridge NSString *)kCFStreamPropertySSLContext];
				SSLProtocol proto;
				SSLGetNegotiatedProtocolVersion(sslContext, &proto);
				[cert setNegotiatedProtocol:proto];
				
				SSLCipherSuite cipher;
				SSLGetNegotiatedCipher(sslContext, &cipher);
				[cert setNegotiatedCipher:cipher];
				
				[[self delegate] HTTPConnection:self didReceiveSecTrust:trust certificate:cert];
			}
		}
		
		NSMutableData *data = [[NSMutableData alloc] initWithCapacity:1024];
		while ([theStream hasBytesAvailable]) {
			uint8_t buf[1024];
			NSUInteger len = [theStream read:buf maxLength:1024];
			[data appendBytes:(const void *)buf length:len];
		}

		[[self delegate] HTTPConnection:self didReceiveData:data];
		
		break;
	}
	default:
		break;
	}
}

@end


#pragma mark -


@implementation CKHTTPConnection (Authentication)

- (CKHTTPAuthenticationChallenge *)currentAuthenticationChallenge {
	return _authenticationChallenge;
}

- (void)_finishCurrentAuthenticationChallenge
{
	_authenticationChallenge = nil;
}

- (void)useCredential:(NSURLCredential *)credential forAuthenticationChallenge:(NSURLAuthenticationChallenge *)challenge
{
	NSParameterAssert(challenge == [self currentAuthenticationChallenge]);
	[self _finishCurrentAuthenticationChallenge];

	// Retry the request, this time with authentication
	// TODO: What if this function fails?
	CFHTTPAuthenticationRef HTTPAuthentication = [(CKHTTPAuthenticationChallenge *)challenge CFHTTPAuthentication];
	CFHTTPMessageApplyCredentials([self HTTPRequest], HTTPAuthentication, (__bridge CFStringRef)[credential user], (__bridge CFStringRef)[credential password], NULL);
	[self start];
}

- (void)continueWithoutCredentialForAuthenticationChallenge:(NSURLAuthenticationChallenge *)challenge
{
	NSParameterAssert(challenge == [self currentAuthenticationChallenge]);
	[self _finishCurrentAuthenticationChallenge];

	// Just return the authentication response to the delegate
	[[self delegate] HTTPConnection:self didReceiveResponse:(NSHTTPURLResponse *)[challenge failureResponse]];
	[[self delegate] HTTPConnectionDidFinishLoading:self];
}

- (void)cancelAuthenticationChallenge:(NSURLAuthenticationChallenge *)challenge
{
	NSParameterAssert(challenge == [self currentAuthenticationChallenge]);
	[self _finishCurrentAuthenticationChallenge];

	// Treat like a -cancel message
	[self cancel];
}

@end


#pragma mark -


@implementation NSURLRequest (CKHTTPURLRequest)

- (CFHTTPMessageRef)makeHTTPMessage
{
	CFHTTPMessageRef result = CFHTTPMessageCreateRequest(NULL, (__bridge CFStringRef)[self HTTPMethod], (__bridge CFURLRef)[self URL], kCFHTTPVersion1_1);
	
	CFHTTPMessageSetHeaderFieldValue(result, (__bridge CFStringRef)@"Accept-Encoding", (__bridge CFStringRef)@"gzip, deflate");
	
	if ([[[self HTTPMethod] uppercaseString] isEqualToString:@"GET"])
		CFHTTPMessageSetHeaderFieldValue(result, (__bridge CFStringRef)@"Connection", (__bridge CFStringRef)@"keep-alive");
	else
		CFHTTPMessageSetHeaderFieldValue(result, (__bridge CFStringRef)@"Connection", (__bridge CFStringRef)@"close");

	for (NSString *hf in [self allHTTPHeaderFields])
		CFHTTPMessageSetHeaderFieldValue(result, (__bridge CFStringRef)hf, (__bridge CFStringRef)[[self allHTTPHeaderFields] objectForKey:hf]);

	NSData *body = [self HTTPBody];
	if (body)
		CFHTTPMessageSetBody(result, (__bridge CFDataRef)body);
	
	return result;
}

@end


@implementation CKHTTPAuthenticationChallenge

/*  Returns nil if the ref is not suitable
 */
- (id)initWithResponse:(CFHTTPMessageRef)response
    proposedCredential:(NSURLCredential *)credential
  previousFailureCount:(NSInteger)failureCount
       failureResponse:(NSHTTPURLResponse *)URLResponse
                sender:(id <NSURLAuthenticationChallengeSender>)sender
{
	NSParameterAssert(response);
	
#warning "Instance variable used while 'self' is not set to the result of [self init]"
	 
	// Try to create an authentication object from the response
	_HTTPAuthentication = CFHTTPAuthenticationCreateFromResponse(NULL, response);
	if (![self CFHTTPAuthentication])
		return nil;
	
	// NSURLAuthenticationChallenge only handles user and password
	if (!CFHTTPAuthenticationIsValid([self CFHTTPAuthentication], NULL))
		return nil;
	
	if (!CFHTTPAuthenticationRequiresUserNameAndPassword([self CFHTTPAuthentication]))
		return nil;
	
	// Fail if we can't retrieve decent protection space info
	CFArrayRef authenticationDomains = CFHTTPAuthenticationCopyDomains([self CFHTTPAuthentication]);
	NSURL *URL = [(__bridge NSArray *)authenticationDomains lastObject];
	CFRelease(authenticationDomains);

	if (!URL || ![URL host])
		return nil;
	
	// Fail for an unsupported authentication method
	CFStringRef authMethod = CFHTTPAuthenticationCopyMethod([self CFHTTPAuthentication]);
	NSString *authenticationMethod;
	if ([(__bridge NSString *)authMethod isEqualToString:(NSString *)kCFHTTPAuthenticationSchemeBasic])
		authenticationMethod = NSURLAuthenticationMethodHTTPBasic;
	else if ([(__bridge NSString *)authMethod isEqualToString:(NSString *)kCFHTTPAuthenticationSchemeDigest])
		authenticationMethod = NSURLAuthenticationMethodHTTPDigest;
	else {
		CFRelease(authMethod);
		 // unsupported authentication scheme
		return nil;
	}
	CFRelease(authMethod);
	
	// Initialise
	CFStringRef realm = CFHTTPAuthenticationCopyRealm([self CFHTTPAuthentication]);
	
	NSURLProtectionSpace *protectionSpace = [[NSURLProtectionSpace alloc] initWithHost:[URL host]
										  port:([URL port] ? [[URL port] intValue] : 80)
									      protocol:[[URL scheme] lowercaseString]
										 realm:(__bridge NSString *)realm
								  authenticationMethod:authenticationMethod];
	CFRelease(realm);

	self = [self initWithProtectionSpace:protectionSpace
		      proposedCredential:credential
		    previousFailureCount:failureCount
			 failureResponse:URLResponse
				   error:nil
				  sender:sender];
	
	return self;
}

- (void)dealloc
{
	CFRelease(_HTTPAuthentication);
}

- (CFHTTPAuthenticationRef)CFHTTPAuthentication {
	return _HTTPAuthentication;
}

@end