/*
 * Copyright 2026 Google LLC
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#import <OCMock/OCMock.h>
#import <XCTest/XCTest.h>

#import <netinet/in.h>
#import <poll.h>
#import <sys/socket.h>
#import <unistd.h>

#import "FirebaseDatabase/Sources/Core/FRepoInfo.h"
#import "FirebaseDatabase/Sources/Realtime/FWebSocketConnection.h"

@interface FWebSocketConnection (Testing)
- (void)closeIfNeverConnected;
#if !TARGET_OS_WATCH
- (void)setWebSocket:(FSRWebSocket *)webSocket;
#endif  // !TARGET_OS_WATCH
@end

#if !TARGET_OS_WATCH

@interface FWebSocketConnection (KeepAliveTesting)
- (void)onClosed;
- (void)resetKeepAlive;
- (void)webSocketDidOpen;
- (void)sendStringToWebSocket:(NSString *)string;
@end

static char workerQueueKey;

@interface FWebSocketKeepAliveSpy : FWebSocketConnection
@property(nonatomic) NSUInteger keepAliveCount;
@property(nonatomic) BOOL sentOnWorkerQueue;
@property(nonatomic, copy) NSString *lastSentString;
@property(nonatomic, copy) void (^onKeepAlive)(void);
@end

@implementation FWebSocketKeepAliveSpy
- (void)sendStringToWebSocket:(NSString *)string {
  self.lastSentString = string;
  self.sentOnWorkerQueue = dispatch_get_specific(&workerQueueKey) != NULL;
  self.keepAliveCount++;
  if (self.onKeepAlive) {
    self.onKeepAlive();
  }
}
@end

@interface FWebSocketConnectionTestDelegate : NSObject <FWebSocketDelegate>
@property(nonatomic) BOOL receivedMessage;
@property(nonatomic) NSInteger disconnectCount;
@property(nonatomic) BOOL wasEverConnected;
@end

@implementation FWebSocketConnectionTestDelegate
- (void)onMessage:(FWebSocketConnection *)fwebSocket withMessage:(NSDictionary *)message {
  self.receivedMessage = YES;
}
- (void)onDisconnect:(FWebSocketConnection *)fwebSocket wasEverConnected:(BOOL)everConnected {
  self.disconnectCount++;
  self.wasEverConnected = everConnected;
}
@end

@interface FWebSocketConnectionTest : XCTestCase
@end

@implementation FWebSocketConnectionTest

- (FWebSocketConnection *)connectionWithDelegate:(FWebSocketConnectionTestDelegate *)delegate {
  return [self connectionToHost:@"foo.firebaseio.com" delegate:delegate];
}

- (FWebSocketConnection *)connectionToHost:(NSString *)host
                                  delegate:(FWebSocketConnectionTestDelegate *)delegate {
  FRepoInfo *info = [[FRepoInfo alloc] initWithHost:host isSecure:YES withNamespace:@"foo"];
  FWebSocketConnection *connection =
      [[FWebSocketConnection alloc] initWith:info
                                    andQueue:dispatch_get_main_queue()
                                 googleAppID:@"1:1234:ios:1234"
                               lastSessionID:nil
                               appCheckToken:nil];
  connection.delegate = delegate;
  return connection;
}

/** Waits up to `timeout` seconds for `fd` to be readable, running the main run loop meanwhile. */
- (BOOL)waitUntilReadable:(int)fd timeout:(NSTimeInterval)timeout {
  NSDate *deadline = [NSDate dateWithTimeIntervalSinceNow:timeout];
  while ([deadline timeIntervalSinceNow] > 0) {
    struct pollfd pollFD = {.fd = fd, .events = POLLIN};
    if (poll(&pollFD, 1, 0) > 0) {
      return YES;
    }
    [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.05]];
  }
  return NO;
}

// The realtime protocol is text only. SocketRocket delivers an NSData for a
// binary frame, which does not respond to the NSString selectors the frame
// handler uses. A server sending such a frame previously crashed the client via
// -[NSData intValue]. It should now be ignored.
- (void)testBinaryFrameIsIgnored {
  FWebSocketConnectionTestDelegate *delegate = [[FWebSocketConnectionTestDelegate alloc] init];
  FWebSocketConnection *connection = [self connectionWithDelegate:delegate];

  NSData *binaryFrame = [@"AB" dataUsingEncoding:NSUTF8StringEncoding];
  XCTAssertNoThrow([connection webSocket:nil didReceiveMessage:(id)binaryFrame]);
  XCTAssertFalse(delegate.receivedMessage);
}

- (void)testEmptyBinaryFrameIsIgnored {
  FWebSocketConnectionTestDelegate *delegate = [[FWebSocketConnectionTestDelegate alloc] init];
  FWebSocketConnection *connection = [self connectionWithDelegate:delegate];

  XCTAssertNoThrow([connection webSocket:nil didReceiveMessage:(id)[NSData data]]);
  XCTAssertFalse(delegate.receivedMessage);
}

// When a connection attempt times out, the disconnect should be reported right
// away, even if the web socket never reports that it closed. The client
// previously waited for that report, which could take hours if the connection
// was stuck in the TLS handshake, and stayed offline until then.
// https://github.com/firebase/firebase-ios-sdk/issues/9682
- (void)testConnectTimeoutReportsDisconnectWithoutWaitingForWebSocket {
  FWebSocketConnectionTestDelegate *delegate = [[FWebSocketConnectionTestDelegate alloc] init];
  FWebSocketConnection *connection = [self connectionWithDelegate:delegate];
  // A web socket that never reports that it closed.
  id webSocket = OCMClassMock([FSRWebSocket class]);
  [connection setWebSocket:webSocket];

  [connection closeIfNeverConnected];

  OCMVerify([webSocket close]);
  // The connection stops listening to the web socket.
  OCMVerify([webSocket setDelegate:nil]);
  XCTAssertEqual(delegate.disconnectCount, 1);
  XCTAssertFalse(delegate.wasEverConnected);

  // The disconnect isn't reported again if the web socket reports it later.
  [connection webSocket:webSocket didCloseWithCode:0 reason:nil wasClean:NO];
  XCTAssertEqual(delegate.disconnectCount, 1);
}

- (void)testConnectTimeoutAfterCloseDoesNotReportDisconnect {
  FWebSocketConnectionTestDelegate *delegate = [[FWebSocketConnectionTestDelegate alloc] init];
  FWebSocketConnection *connection = [self connectionWithDelegate:delegate];
  id webSocket = OCMClassMock([FSRWebSocket class]);
  [connection setWebSocket:webSocket];

  [connection close];
  [connection closeIfNeverConnected];

  XCTAssertEqual(delegate.disconnectCount, 0);
}

// When a connection attempt times out, the web socket should close the
// connection even if it's stuck in the TLS handshake. It previously waited to
// send its handshake request first, which kept the connection open until the
// OS gave up on it, possibly hours later.
// https://github.com/firebase/firebase-ios-sdk/issues/9682
- (void)testConnectTimeoutClosesConnectionStuckInTLSHandshake {
  // Stop at the first failure, so that a failed wait doesn't lead to a blocking
  // accept or read.
  self.continueAfterFailure = NO;
  // A local server that accepts the connection but never answers the TLS
  // handshake.
  int server = socket(AF_INET, SOCK_STREAM, 0);
  XCTAssertGreaterThanOrEqual(server, 0);
  [self addTeardownBlock:^{
    close(server);
  }];
  struct sockaddr_in address = {0};
  address.sin_len = sizeof(address);
  address.sin_family = AF_INET;
  address.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  socklen_t addressLength = sizeof(address);
  XCTAssertEqual(bind(server, (struct sockaddr *)&address, addressLength), 0);
  XCTAssertEqual(listen(server, 1), 0);
  XCTAssertEqual(getsockname(server, (struct sockaddr *)&address, &addressLength), 0);
  NSString *host = [NSString stringWithFormat:@"127.0.0.1:%d", ntohs(address.sin_port)];

  FWebSocketConnectionTestDelegate *delegate = [[FWebSocketConnectionTestDelegate alloc] init];
  FWebSocketConnection *connection = [self connectionToHost:host delegate:delegate];
  [connection open];

  XCTAssertTrue([self waitUntilReadable:server timeout:5]);
  int client = accept(server, NULL, NULL);
  XCTAssertGreaterThanOrEqual(client, 0);
  [self addTeardownBlock:^{
    close(client);
  }];
  // Wait for the TLS ClientHello, then give the web socket time to queue its
  // handshake request.
  char buffer[4096];
  XCTAssertTrue([self waitUntilReadable:client timeout:5]);
  XCTAssertGreaterThan(read(client, buffer, sizeof(buffer)), 0);
  [[NSRunLoop currentRunLoop] runUntilDate:[NSDate dateWithTimeIntervalSinceNow:0.5]];

  [connection closeIfNeverConnected];
  XCTAssertEqual(delegate.disconnectCount, 1);

  // Read until the connection is closed. The client may send a TLS alert first.
  BOOL closed = NO;
  while (!closed && [self waitUntilReadable:client timeout:5]) {
    closed = read(client, buffer, sizeof(buffer)) <= 0;
  }
  XCTAssertTrue(closed);
  XCTAssertEqual(delegate.disconnectCount, 1);
}

- (FWebSocketKeepAliveSpy *)keepAliveConnectionOnQueue:(dispatch_queue_t)queue {
  FRepoInfo *info = [[FRepoInfo alloc] initWithHost:@"foo.firebaseio.com"
                                           isSecure:YES
                                      withNamespace:@"foo"];
  dispatch_queue_set_specific(queue, &workerQueueKey, &workerQueueKey, NULL);
  return [[FWebSocketKeepAliveSpy alloc] initWith:info
                                         andQueue:queue
                                      googleAppID:@"1:1234:ios:1234"
                                    lastSessionID:nil
                                    appCheckToken:nil];
}

- (void)testKeepAliveRunsOnWorkerQueueAndResetsInline {
  dispatch_queue_t queue = dispatch_queue_create("keepalive.test", DISPATCH_QUEUE_SERIAL);
  FWebSocketKeepAliveSpy *connection = [self keepAliveConnectionOnQueue:queue];
  XCTestExpectation *sent = [self expectationWithDescription:@"worker keepalive"];
  dispatch_async(queue, ^{
    [connection webSocketDidOpen];
    [connection setValue:[NSDate distantPast] forKey:@"keepAliveFireDate"];
    [connection resetKeepAlive];
    NSDate *fireDate = [connection valueForKey:@"keepAliveFireDate"];
    XCTAssertGreaterThan([fireDate timeIntervalSinceNow], 40.0);
    __weak FWebSocketKeepAliveSpy *weakConnection = connection;
    connection.onKeepAlive = ^{
      FWebSocketKeepAliveSpy *activeConnection = weakConnection;
      XCTAssertNotNil(activeConnection);
      XCTAssertEqualObjects(activeConnection.lastSentString, @"0");
      XCTAssertTrue(activeConnection.sentOnWorkerQueue);
      [activeConnection onClosed];
      activeConnection.onKeepAlive = nil;
      [sent fulfill];
    };
    dispatch_source_t timer = [connection valueForKey:@"keepAlive"];
    dispatch_source_set_timer(timer, DISPATCH_TIME_NOW, DISPATCH_TIME_FOREVER, 0);
  });
  [self waitForExpectations:@[ sent ] timeout:2];
}

- (void)testCloseCancelsKeepAliveAndCannotReinstallIt {
  dispatch_queue_t queue = dispatch_queue_create("keepalive.close.test", DISPATCH_QUEUE_SERIAL);
  FWebSocketKeepAliveSpy *connection = [self keepAliveConnectionOnQueue:queue];
  XCTestExpectation *closed = [self expectationWithDescription:@"no keepalive after close"];
  dispatch_async(queue, ^{
    [connection webSocketDidOpen];
    dispatch_source_t timer = [connection valueForKey:@"keepAlive"];
    dispatch_source_set_timer(timer, dispatch_time(DISPATCH_TIME_NOW, 100 * NSEC_PER_MSEC),
                              DISPATCH_TIME_FOREVER, 0);
    [connection close];
    [connection resetKeepAlive];
    [connection webSocketDidOpen];
    XCTAssertNil([connection valueForKey:@"keepAlive"]);
    dispatch_after(dispatch_time(DISPATCH_TIME_NOW, 200 * NSEC_PER_MSEC), queue, ^{
      XCTAssertEqual(connection.keepAliveCount, 0U);
      [closed fulfill];
    });
  });
  [self waitForExpectations:@[ closed ] timeout:2];
}

@end

#endif  // !TARGET_OS_WATCH

// Records the disconnects that a connection reports.
@interface FWebSocketConnectionDisconnectRecorder : NSObject <FWebSocketDelegate>
@property(nonatomic) NSInteger disconnectCount;
@property(nonatomic) BOOL wasEverConnected;
@property(nonatomic) XCTestExpectation *disconnectExpectation;
@end

@implementation FWebSocketConnectionDisconnectRecorder
- (void)onMessage:(FWebSocketConnection *)fwebSocket withMessage:(NSDictionary *)message {
}
- (void)onDisconnect:(FWebSocketConnection *)fwebSocket wasEverConnected:(BOOL)everConnected {
  self.disconnectCount++;
  self.wasEverConnected = everConnected;
  [self.disconnectExpectation fulfill];
}
@end

// A connection URL that NSURL can't parse used to abort the app in
// assert(request.URL) in -[FSRWebSocket initWithURLRequest:...]. These tests
// also run on watchOS, which uses an NSURLSession web socket task instead.
@interface FWebSocketConnectionInvalidURLTest : XCTestCase
@end

@implementation FWebSocketConnectionInvalidURLTest

- (FWebSocketConnection *)connectionWithInvalidURLAndDelegate:
    (FWebSocketConnectionDisconnectRecorder *)delegate {
  // NSURL rejects a host with a space.
  FRepoInfo *info = [[FRepoInfo alloc] initWithHost:@"bad host" isSecure:YES withNamespace:@"foo"];
  XCTAssertNil([NSURL URLWithString:[info connectionURL]]);
  FWebSocketConnection *connection =
      [[FWebSocketConnection alloc] initWith:info
                                    andQueue:dispatch_get_main_queue()
                                 googleAppID:@"1:1234:ios:1234"
                               lastSessionID:nil
                               appCheckToken:nil];
  connection.delegate = delegate;
  return connection;
}

- (void)testInvalidURLReportsFailedConnectionAttempt {
  FWebSocketConnectionDisconnectRecorder *delegate =
      [[FWebSocketConnectionDisconnectRecorder alloc] init];
  delegate.disconnectExpectation = [self expectationWithDescription:@"disconnect"];
  FWebSocketConnection *connection = [self connectionWithInvalidURLAndDelegate:delegate];

  [connection open];

  // Reported asynchronously, as a web socket would, so open returns before the
  // delegate is called.
  XCTAssertEqual(delegate.disconnectCount, 0);

  // Reported right away rather than after the 30 second connect timeout, so
  // the connection falls back to the configured host and retries with backoff.
  [self waitForExpectations:@[ delegate.disconnectExpectation ] timeout:5];
  XCTAssertEqual(delegate.disconnectCount, 1);
  XCTAssertFalse(delegate.wasEverConnected);

  // The connect timeout doesn't report it again.
  [connection closeIfNeverConnected];
  XCTAssertEqual(delegate.disconnectCount, 1);
}

// As with a web socket, closing the connection before the failure is reported
// suppresses the report.
- (void)testInvalidURLFailureIsNotReportedAfterClose {
  FWebSocketConnectionDisconnectRecorder *delegate =
      [[FWebSocketConnectionDisconnectRecorder alloc] init];
  FWebSocketConnection *connection = [self connectionWithInvalidURLAndDelegate:delegate];

  [connection open];
  [connection close];

  // The report is dispatched to the main queue before this block.
  XCTestExpectation *reportDispatched = [self expectationWithDescription:@"report dispatched"];
  dispatch_async(dispatch_get_main_queue(), ^{
    [reportDispatched fulfill];
  });
  [self waitForExpectations:@[ reportDispatched ] timeout:5];
  XCTAssertEqual(delegate.disconnectCount, 0);
}

@end