// Copyright (c) 2013, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.
//
// VMOptions=
// VMOptions=--short_socket_read
// VMOptions=--short_socket_write
// VMOptions=--short_socket_read --short_socket_write
//
// OtherResources=certificates/server_chain.pem
// OtherResources=certificates/server_key.pem
// OtherResources=certificates/trusted_certs.pem

import "dart:async";
import "dart:convert";
import "dart:io";
import "dart:typed_data";

import "package:crypto/crypto.dart";
import "package:expect/async_helper.dart";
import "package:expect/expect.dart";

const WEB_SOCKET_GUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";

const String HOST_NAME = 'localhost';

String localFile(path) => Platform.script.resolve(path).toFilePath();

SecurityContext serverContext = new SecurityContext()
  ..useCertificateChain(localFile('certificates/server_chain.pem'))
  ..usePrivateKey(
    localFile('certificates/server_key.pem'),
    password: 'dartdart',
  );

SecurityContext clientContext = new SecurityContext()
  ..setTrustedCertificates(localFile('certificates/trusted_certs.pem'));

/**
 * A SecurityConfiguration lets us run the tests over HTTP or HTTPS.
 */
class SecurityConfiguration {
  final bool secure;

  SecurityConfiguration({required this.secure});

  Future<HttpServer> createServer({int backlog = 0}) => secure
      ? HttpServer.bindSecure(HOST_NAME, 0, serverContext, backlog: backlog)
      : HttpServer.bind(HOST_NAME, 0, backlog: backlog);

  Future<WebSocket> createClient(
    int port, {
    String? user,
    Map<String, Object>? headers,
    String? customUserAgent,
  }) => WebSocket.connect(
    '${secure ? "wss" : "ws"}://${user is Null ? "" : "$user@"}$HOST_NAME:$port/',
    headers: headers,
    customClient: secure
        ? (HttpClient(context: clientContext)..userAgent = customUserAgent)
        : null,
  );

  checkCloseStatus(webSocket, closeStatus, closeReason) {
    Expect.equals(
      closeStatus == null ? WebSocketStatus.noStatusReceived : closeStatus,
      webSocket.closeCode,
    );
    Expect.equals(
      closeReason == null ? "" : closeReason,
      webSocket.closeReason,
    );
  }

  void testRequestResponseClientCloses(
    int totalConnections,
    int? closeStatus,
    String? closeReason,
    int numberOfMessages,
  ) {
    assert(numberOfMessages >= 1);

    asyncStart();
    createServer().then((server) {
      server
          .transform(new WebSocketTransformer())
          .listen(
            (webSocket) {
              asyncStart();
              webSocket.listen(
                webSocket.add,
                onDone: () {
                  checkCloseStatus(webSocket, closeStatus, closeReason);
                  asyncEnd();
                },
              );
            },
            onDone: () {
              asyncEnd();
            },
          );

      int closeCount = 0;
      String messageText = "Hello, world!";
      for (int i = 0; i < totalConnections; i++) {
        asyncStart();
        createClient(server.port).then((webSocket) {
          webSocket.add(messageText);
          webSocket.listen(
            (message) {
              numberOfMessages--;
              Expect.equals(messageText, message);

              if (numberOfMessages > 0) {
                webSocket.add(message);
              } else {
                webSocket.close(closeStatus, closeReason);
              }
            },
            onDone: () {
              checkCloseStatus(webSocket, closeStatus, closeReason);
              closeCount++;
              if (closeCount == totalConnections) {
                server.close();
              }
              asyncEnd();
            },
          );
        });
      }
    });
  }

  void testRequestResponseServerCloses(
    int totalConnections,
    int? closeStatus,
    String? closeReason,
  ) {
    createServer().then((server) {
      int closeCount = 0;
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        String messageText = "Hello, world!";
        int messageCount = 0;
        webSocket.listen(
          (message) {
            messageCount++;
            if (messageCount < 10) {
              Expect.equals(messageText, message);
              webSocket.add(message);
            } else {
              webSocket.close(closeStatus, closeReason);
            }
          },
          onDone: () {
            checkCloseStatus(webSocket, closeStatus, closeReason);
            closeCount++;
            if (closeCount == totalConnections) {
              server.close();
            }
          },
        );
        webSocket.add(messageText);
      });

      for (int i = 0; i < totalConnections; i++) {
        createClient(server.port).then((webSocket) {
          webSocket.listen(
            webSocket.add,
            onDone: () {
              checkCloseStatus(webSocket, closeStatus, closeReason);
            },
          );
        });
      }
    });
  }

  void testMessageLength(int messageLength) {
    createServer().then((server) {
      Uint8List originalMessage = new Uint8List(messageLength);
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        webSocket.listen((message) {
          Expect.listEquals(originalMessage, message);
          webSocket.add(message);
        });
      });

      createClient(server.port).then((webSocket) {
        webSocket.listen((message) {
          Expect.listEquals(originalMessage, message);
          webSocket.close();
        }, onDone: server.close);
        webSocket.add(originalMessage);
      });
    });
  }

  void testCloseNoListen() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        server.close();
        webSocket.close();
      });

      createClient(server.port).then((webSocket) {
        webSocket.close();
      });
    });
  }

  void testCancelThenClose() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        webSocket.listen(null).cancel();
        webSocket.close();
        server.close();
      });

      createClient(server.port).then((webSocket) {
        webSocket.close();
      });
    });
  }

  void testCloseThenCancel() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        var subscription = webSocket.listen(null);
        webSocket.close();
        subscription.cancel();
        server.close();
      });

      createClient(server.port).then((webSocket) {
        webSocket.close();
      });
    });
  }

  void testListenAfterClose() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        server.close();
        webSocket.close();
        Expect.throws(() => webSocket.drain());
      });

      createClient(server.port).then((webSocket) {
        webSocket.close();
        Expect.throws(() => webSocket.drain());
      });
    });
  }

  void testDoubleCloseClient() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        server.close();
        webSocket.listen((_) {}, onDone: webSocket.close);
      });

      createClient(server.port).then((webSocket) {
        webSocket.listen((_) {}, onDone: webSocket.close);
        webSocket.close();
      });
    });
  }

  void testDoubleCloseServer() {
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        server.close();
        webSocket.listen((_) {}, onDone: webSocket.close);
        webSocket.close();
      });

      createClient(server.port).then((webSocket) {
        webSocket.listen((_) {}, onDone: webSocket.close);
      });
    });
  }

  void testImmediateCloseServer() {
    createServer().then((server) {
      server.listen((request) {
        WebSocketTransformer.upgrade(request).then((webSocket) {
          webSocket.listen((_) {
            Expect.fail("Unexpected message");
          }, onDone: server.close);
          webSocket.close();
        });
      });

      createClient(server.port).then((webSocket) {
        webSocket.listen((_) {
          Expect.fail("Unexpected message");
        }, onDone: webSocket.close);
      });
    });
  }

  void testImmediateCloseClient() {
    createServer().then((server) {
      server.listen((request) {
        WebSocketTransformer.upgrade(request).then((webSocket) {
          webSocket.listen(
            (_) {
              Expect.fail("Unexpected message");
            },
            onDone: () {
              server.close();
              webSocket.close();
            },
          );
        });
      });

      createClient(server.port).then((webSocket) {
        webSocket.listen((_) {
          Expect.fail("Unexpected message");
        }, onDone: webSocket.close);
        webSocket.close();
      });
    });
  }

  void testNoUpgrade() {
    createServer().then((server) {
      // Create a server which always responds with notFound.
      server.listen((request) {
        request.response.statusCode = HttpStatus.notFound;
        request.response.close();
      });

      Future<WebSocket?>.value(createClient(server.port)).catchError((error) {
        Expect.equals(HttpStatus.notFound, error.httpStatusCode);
        server.close();
      });
    });
  }

  void testUsePOST() {
    asyncStart();
    createServer().then((server) {
      server
          .transform(new WebSocketTransformer())
          .listen(
            (webSocket) {
              Expect.fail("No connection expected");
            },
            onError: (e) {
              asyncEnd();
            },
          );

      final client = HttpClient(context: secure ? clientContext : null);
      client
          .postUrl(
            Uri.parse(
              "${secure ? 'https:' : 'http:'}//$HOST_NAME:${server.port}/",
            ),
          )
          .then((request) => request.close())
          .then((response) {
            Expect.equals(HttpStatus.badRequest, response.statusCode);
            client.close();
            server.close();
          });
    });
  }

  void testConnections(
    int totalConnections,
    int closeStatus,
    String closeReason,
  ) {
    createServer().then((server) {
      int closeCount = 0;
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        String messageText = "Hello, world!";
        int messageCount = 0;
        webSocket.listen(
          (message) {
            messageCount++;
            if (messageCount < 10) {
              Expect.equals(messageText, message);
              webSocket.add(message);
            } else {
              webSocket.close(closeStatus, closeReason);
            }
          },
          onDone: () {
            checkCloseStatus(webSocket, closeStatus, closeReason);
            closeCount++;
            if (closeCount == totalConnections) {
              server.close();
            }
          },
        );
        webSocket.add(messageText);
      });

      void webSocketConnection() {
        bool onopenCalled = false;
        int onmessageCalled = 0;
        bool oncloseCalled = false;

        createClient(server.port).then((webSocket) {
          Expect.isFalse(onopenCalled);
          Expect.equals(0, onmessageCalled);
          Expect.isFalse(oncloseCalled);
          onopenCalled = true;
          Expect.equals(WebSocket.open, webSocket.readyState);
          webSocket.listen(
            (message) {
              onmessageCalled++;
              Expect.isTrue(onopenCalled);
              Expect.isFalse(oncloseCalled);
              Expect.equals(WebSocket.open, webSocket.readyState);
              webSocket.add(message);
            },
            onDone: () {
              Expect.isTrue(onopenCalled);
              Expect.equals(10, onmessageCalled);
              Expect.isFalse(oncloseCalled);
              oncloseCalled = true;
              Expect.equals(3002, webSocket.closeCode);
              Expect.equals("Got tired", webSocket.closeReason);
              Expect.equals(WebSocket.closed, webSocket.readyState);
            },
          );
        });
      }

      for (int i = 0; i < totalConnections; i++) {
        webSocketConnection();
      }
    });
  }

  testIndividualUpgrade(int connections) {
    asyncStart();
    createServer().then((server) {
      server.listen((request) {
        if (WebSocketTransformer.isUpgradeRequest(request)) {
          WebSocketTransformer.upgrade(request).then((webSocket) {
            webSocket.listen((_) {
              webSocket.close();
            });
            webSocket.add("Hello");
          });
        } else {
          Expect.isFalse(WebSocketTransformer.isUpgradeRequest(request));
          request.response.statusCode = HttpStatus.ok;
          request.response.close();
        }
      });

      var futures = <Future>[];

      var wsProtocol = '${secure ? "wss" : "ws"}';
      var baseWsUrl = '$wsProtocol://$HOST_NAME:${server.port}/';
      var httpProtocol = '${secure ? "https" : "http"}';
      var baseHttpUrl = '$httpProtocol://$HOST_NAME:${server.port}/';
      final client = HttpClient(context: secure ? clientContext : null);

      for (int i = 0; i < connections; i++) {
        var completer = new Completer();
        futures.add(completer.future);
        createClient(server.port).then((websocket) {
          websocket.listen((_) {
            websocket.close();
          }, onDone: completer.complete);
        });

        futures.add(
          client
              .openUrl("GET", Uri.parse('${baseHttpUrl}'))
              .then((request) => request.close())
              .then((response) {
                response.listen((_) {});
                Expect.equals(HttpStatus.ok, response.statusCode);
              }),
        );
      }

      Future.wait(futures).then((_) {
        server.close();
        client.close();
        asyncEnd();
      });
    });
  }

  testFromUpgradedSocket() {
    asyncStart();
    createServer().then((server) {
      server.listen((request) {
        Expect.equals(
          'Upgrade',
          request.headers.value(HttpHeaders.connectionHeader),
        );
        Expect.equals(
          'websocket',
          request.headers.value(HttpHeaders.upgradeHeader),
        );

        var key = request.headers.value('Sec-WebSocket-Key');
        var digest = sha1.convert("$key$WEB_SOCKET_GUID".codeUnits);
        var accept = base64.encode(digest.bytes);
        request.response
          ..statusCode = HttpStatus.switchingProtocols
          ..headers.add(HttpHeaders.connectionHeader, "Upgrade")
          ..headers.add(HttpHeaders.upgradeHeader, "websocket")
          ..headers.add("Sec-WebSocket-Accept", accept);
        request.response.contentLength = 0;
        request.response
            .detachSocket()
            .then((socket) {
              return new WebSocket.fromUpgradedSocket(socket, serverSide: true);
            })
            .then((websocket) {
              websocket.add("Hello");
              websocket.close();
              asyncEnd();
            });
      });

      createClient(server.port)
          .then((websocket) {
            return websocket.listen((message) {
              Expect.equals("Hello", message);
              websocket.close();
            }).asFuture();
          })
          .then((_) => server.close());
    });
  }

  void testAdditionalHeaders() {
    asyncStart();
    createServer().then((server) {
      server.listen((request) {
        Expect.isTrue(WebSocketTransformer.isUpgradeRequest(request));
        Expect.equals('my-value', request.headers['My-Header']![0]);
        var header = request.headers['My-Header-Multiple']!;
        Expect.equals(1, header.length);
        Expect.equals('my-value-1, my-value-2', header[0]);
        WebSocketTransformer.upgrade(request).then((webSocket) {
          webSocket.listen((_) {
            webSocket.close();
          });
          webSocket.add("Hello");
        });
      });

      var headers = {
        'My-Header': 'my-value',
        'My-Header-Multiple': ['my-value-1', 'my-value-2'],
      };
      createClient(server.port, headers: headers)
          .then((websocket) {
            return websocket.listen((message) {
              Expect.equals("Hello", message);
              websocket.close();
            }).asFuture();
          })
          .then((_) {
            server.close();
            asyncEnd();
          });
    });
  }

  void testBasicAuthentication() {
    var userInfo = 'user:password';

    asyncStart();
    asyncStart();
    createServer().then((server) {
      server.listen((request) {
        Expect.isTrue(WebSocketTransformer.isUpgradeRequest(request));
        String auth = base64.encode(utf8.encode(userInfo));
        Expect.equals('Basic $auth', request.headers['Authorization']![0]);
        Expect.equals(1, request.headers['Authorization']!.length);
        WebSocketTransformer.upgrade(request).then((webSocket) {
          webSocket.listen(
            (_) {
              throw 'Unexpected';
            },
            onDone: () {
              asyncEnd();
            },
          );
          webSocket.add("Hello");
        });
      });

      createClient(server.port, user: userInfo)
          .then((websocket) {
            return websocket.listen((message) {
              Expect.equals("Hello", message);
              websocket.close();
            }).asFuture();
          })
          .then((_) {
            return server.close();
          })
          .whenComplete(() {
            asyncEnd();
          });
    });
  }

  void testShouldSetUserAgent() {
    asyncStart();
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        Expect.equals('Custom User Agent', WebSocket.userAgent);
        server.close();
        webSocket.close();
        asyncEnd();
      });

      WebSocket.userAgent = 'Custom User Agent';
      createClient(server.port).then((webSocket) {
        webSocket.close();
      });
    });
  }

  void testStaticClientUserAgentStaysTheSame() {
    asyncStart();
    createServer().then((server) {
      server.transform(new WebSocketTransformer()).listen((webSocket) {
        Expect.equals('Custom User Agent', WebSocket.userAgent);
        server.close();
        webSocket.close();
        asyncEnd();
      });
      // Next line should take no effect on custom user agent value provided
      WebSocket.userAgent = 'Custom User Agent';
      createClient(server.port, customUserAgent: 'New User Agent').then((
        webSocket,
      ) {
        webSocket.close();
      });
    });
  }

  void runTests() {
    testRequestResponseClientCloses(2, null, null, 1);
    testRequestResponseClientCloses(2, 3001, null, 2);
    testRequestResponseClientCloses(2, 3002, "Got tired", 3);
    testRequestResponseServerCloses(2, null, null);
    testRequestResponseServerCloses(2, 3001, null);
    testRequestResponseServerCloses(2, 3002, "Got tired");
    testMessageLength(125);
    testMessageLength(126);
    testMessageLength(127);
    testMessageLength(65535);
    testMessageLength(65536);
    testCloseNoListen();
    testCancelThenClose();
    testCloseThenCancel();
    testListenAfterClose();
    testDoubleCloseClient();
    testDoubleCloseServer();
    testImmediateCloseServer();
    testImmediateCloseClient();
    testNoUpgrade();
    testUsePOST();
    testConnections(10, 3002, "Got tired");
    testIndividualUpgrade(5);
    testFromUpgradedSocket();
    testAdditionalHeaders();
    testBasicAuthentication();
    testShouldSetUserAgent();
    testStaticClientUserAgentStaysTheSame();
  }
}

main() {
  new SecurityConfiguration(secure: false).runTests();
  new SecurityConfiguration(secure: true).runTests();
}