authentication overhaul

.gitignore

@@ -109,3 +109,4 @@ dist
 # TernJS port file
 .tern-port
 .DS_Store
+puppeteer-data/

README.md

@@ -19,9 +19,26 @@ yarn build
 
 ```
 yarn install
-node dist/example/testApp.js
+# if you already have an access token and refresh token
+DEBUG_DEPTH=5 DEBUG=* TOS_ACCESS_TOKEN=<access_token> TOS_REFRESH_TOKEN=<refresh_token>  node dist/example/testApp.js
+# if you don't have an access token and refresh token, you can use your username and password
+# this will launch a browser to authenticate and then save the access token and refresh token to .env
+DEBUG_DEPTH=5 DEBUG=* TOS_USERNAME=<username> TOS_PASSWORD=<password> node dist/example/testApp.js
 ```
 
+## Authentication flow
+
+There seems to be currently two ways to authenticate:
+
+1. From scratch with username and password:
+   a. Send a message with the `login/schwab` service including the `authCode` obtained from the browser oauth flow at `trade.thinkorswim.com`;
+   b. This will return a `token` and `refreshToken` which should be saved for future use;
+   c. The `authCode` is single use and cannot be used again once exchanged for a token.
+2. From a previously obtained `token`
+   a. Send a message with the `login` service including the `token` returned from the `login/schwab` response message (step 1 above).
+   b. This will return the same `token`, weirdly, and a `refreshToken`, which should be saved for future use
+   c. The token is valid for 24 hours.
+
 # Supported APIs
 
 - ✅ Authentication via access token

src/client/mockWsJsonClient.ts

@@ -34,17 +34,26 @@ import { MarketDepthResponse } from "./services/marketDepthMessageHandler.js";
 import { GetWatchlistResponse } from "./services/getWatchlistMessageHandler.js";
 
 export default class MockWsJsonClient implements WsJsonClient {
+  authenticateWithAuthCode(
+    _authCode: string
+  ): Promise<RawLoginResponseBody | null> {
+    throw new Error("Method not implemented.");
+  }
+
+  authenticateWithAccessToken(_: {
+    accessToken: string;
+    refreshToken: string;
+  }): Promise<RawLoginResponseBody | null> {
+    throw new Error("Method not implemented.");
+  }
+
   async *accountPositions(_: string): AsyncIterable<PositionsResponse> {
     return yield {
       service: "positions",
       positions: [],
     };
   }
 
-  authenticate(): Promise<RawLoginResponseBody | null> {
-    return Promise.resolve(null);
-  }
-
   cancelAlert(_: number): Promise<CancelAlertResponse> {
     return Promise.resolve({
       service: "alerts/cancel",

src/client/realWsJsonClient.ts

@@ -8,6 +8,7 @@ import {
 import {
   Constructor,
   debugLog,
+  ensure,
   findByTypeOrThrow,
   throwError,
 } from "./util.js";
@@ -86,6 +87,7 @@ import GetWatchlistMessageHandler, {
   GetWatchlistResponse,
 } from "./services/getWatchlistMessageHandler.js";
 import SchwabLoginMessageHandler from "./services/schwabLoginMessageHandler.js";
+import { writeFileSync, readFileSync, existsSync } from "node:fs";
 
 export const CONNECTION_REQUEST_MESSAGE = {
   ver: "27.*.*",
@@ -130,11 +132,11 @@ export default class RealWsJsonClient implements WsJsonClient {
   private buffer = new BufferedIterator<ParsedWebSocketResponse>();
   private iterator = new MulticastIterator(this.buffer);
   private state = ChannelState.DISCONNECTED;
-  private authCode?: string;
-  // @ts-expect-error
-  private accessToken?: string;
-  // @ts-expect-error
-  private refreshToken?: string;
+  private credentials: {
+    authCode?: string;
+    accessToken?: string;
+    refreshToken?: string;
+  } = {};
 
   constructor(
     private readonly socket = new WebSocket(
@@ -161,8 +163,28 @@ export default class RealWsJsonClient implements WsJsonClient {
     )
   ) {}
 
-  async authenticate(authCode: string): Promise<RawLoginResponseBody | null> {
-    this.authCode = authCode;
+  async authenticateWithAccessToken({
+    accessToken,
+    refreshToken,
+  }: {
+    accessToken: string;
+    refreshToken: string;
+  }): Promise<RawLoginResponseBody | null> {
+    ensure(accessToken, "access token is required");
+    ensure(refreshToken, "refresh token is required");
+    this.credentials = { accessToken, refreshToken };
+    return await this.handshake();
+  }
+
+  async authenticateWithAuthCode(
+    authCode: string
+  ): Promise<RawLoginResponseBody | null> {
+    ensure(authCode, "auth code is required");
+    this.credentials = { authCode };
+    return await this.handshake();
+  }
+
+  private async handshake(): Promise<RawLoginResponseBody | null> {
     const { state } = this;
     switch (state) {
       case ChannelState.DISCONNECTED:
@@ -178,7 +200,6 @@ export default class RealWsJsonClient implements WsJsonClient {
         return Promise.reject("Illegal state, ws connection failed previously");
     }
   }
-
   private doConnect(): Promise<RawLoginResponseBody> {
     const { socket } = this;
     return new Promise((resolve, reject) => {
@@ -198,18 +219,11 @@ export default class RealWsJsonClient implements WsJsonClient {
     resolve: (value: RawLoginResponseBody) => void,
     reject: (reason?: string) => void
   ) {
-    const { responseParser, buffer, authCode } = this;
+    const { responseParser, buffer } = this;
     const message = JSON.parse(data) as WsJsonRawMessage;
     logger("⬅️\treceived %O", message);
     if (isConnectionResponse(message)) {
-      const handler = findByTypeOrThrow(
-        messageHandlers,
-        SchwabLoginMessageHandler
-      );
-      if (!authCode) {
-        throwError("auth code is required, cannot authenticate");
-      }
-      this.sendMessage(handler.buildRequest(authCode));
+      this.authenticate();
     } else if (isLoginResponse(message)) {
       this.handleLoginResponse(message, resolve, reject);
     } else if (isSchwabLoginResponse(message)) {
@@ -222,6 +236,26 @@ export default class RealWsJsonClient implements WsJsonClient {
     }
   }
 
+  private authenticate() {
+    const {
+      credentials: { authCode, accessToken },
+    } = this;
+    if (accessToken) {
+      // if we already have an access token, we can just authenticate with it
+      const handler = findByTypeOrThrow(messageHandlers, LoginMessageHandler);
+      this.sendMessage(handler.buildRequest(accessToken));
+    } else if (authCode) {
+      // exchange the auth code for an access token
+      const handler = findByTypeOrThrow(
+        messageHandlers,
+        SchwabLoginMessageHandler
+      );
+      this.sendMessage(handler.buildRequest(authCode));
+    } else {
+      throwError("no credentials provided, cannot authenticate");
+    }
+  }
+
   isConnected(): boolean {
     const { socket, state } = this;
     return socket !== null && state === ChannelState.CONNECTED;
@@ -364,6 +398,36 @@ export default class RealWsJsonClient implements WsJsonClient {
     this.socket?.send(msg);
   }
 
+  private updateDotEnvCredentials() {
+    const {
+      credentials: { accessToken, refreshToken },
+    } = this;
+    const envPath = ".env";
+    let envContent = "";
+    if (existsSync(envPath)) {
+      envContent = readFileSync(envPath, "utf-8");
+      // Remove any existing token lines
+      envContent = envContent
+        .split("\n")
+        .filter(
+          (line) =>
+            !line.startsWith("TOS_ACCESS_TOKEN=") &&
+            !line.startsWith("TOS_REFRESH_TOKEN=")
+        )
+        .join("\n");
+    }
+
+    // Append the new token values
+    const tokenLines = [
+      `TOS_ACCESS_TOKEN=${accessToken}`,
+      `TOS_REFRESH_TOKEN=${refreshToken}`,
+    ].join("\n");
+
+    // Ensure there's a newline between existing content and new tokens
+    const newContent = envContent.trim() + "\n" + tokenLines + "\n";
+    writeFileSync(envPath, newContent);
+  }
+
   private handleSchwabLoginResponse(
     message: RawLoginResponse,
     resolve: (value: RawLoginResponseBody) => void,
@@ -378,8 +442,11 @@ export default class RealWsJsonClient implements WsJsonClient {
     if (loginResponse.authenticated) {
       this.state = ChannelState.CONNECTED;
       logger("Schwab login successful, token=%s", body.token);
-      this.accessToken = body.token;
-      this.refreshToken = loginResponse.refreshToken;
+      this.credentials.accessToken = body.token;
+      if (loginResponse.refreshToken) {
+        this.credentials.refreshToken = loginResponse.refreshToken;
+      }
+      this.updateDotEnvCredentials();
       resolve(body);
     } else {
       this.state = ChannelState.ERROR;
@@ -398,6 +465,7 @@ export default class RealWsJsonClient implements WsJsonClient {
     const [{ body }] = message.payload;
     if (loginResponse.successful) {
       this.state = ChannelState.CONNECTED;
+      this.updateDotEnvCredentials();
       resolve(body);
     } else {
       this.state = ChannelState.ERROR;

src/client/util.ts

@@ -12,6 +12,13 @@ export function throwError(msg: string): never {
   throw new Error(msg);
 }
 
+export function ensure<T>(value: T, msg: string): T {
+  if (!value) {
+    throwError(msg);
+  }
+  return value;
+}
+
 export function positionNetQuantity(position: AccountPosition): number {
   return position.longQuantity - position.shortQuantity;
 }

src/client/wsJsonClient.ts

@@ -34,7 +34,14 @@ import { GetWatchlistResponse } from "./services/getWatchlistMessageHandler.js";
 import { Disposable } from "../server/disposable.js";
 
 export interface WsJsonClient extends Disposable {
-  authenticate(accessToken: string): Promise<RawLoginResponseBody | null>;
+  authenticateWithAuthCode(
+    authCode: string
+  ): Promise<RawLoginResponseBody | null>;
+
+  authenticateWithAccessToken(args: {
+    accessToken: string;
+    refreshToken: string;
+  }): Promise<RawLoginResponseBody | null>;
 
   isConnected(): boolean;
 

src/client/wsJsonClientProxy.ts

@@ -42,7 +42,8 @@ import { Observable, BufferedIterator, MulticastIterator } from "obgen";
 const logger = debug("wsClientProxy");
 
 export const ALL_REQUESTS = [
-  "authenticate",
+  "authenticateWithAuthCode",
+  "authenticateWithAccessToken",
   "optionChainQuotes",
   "disconnect",
   "quotes",
@@ -86,8 +87,22 @@ export default class WsJsonClientProxy implements WsJsonClient {
     private readonly options?: any
   ) {}
 
-  async authenticate(
-    accessToken: string
+  authenticateWithAuthCode(
+    authCode: string
+  ): Promise<RawLoginResponseBody | null> {
+    return this.authenticate("authenticateWithAuthCode", authCode);
+  }
+
+  authenticateWithAccessToken(args: {
+    accessToken: string;
+    refreshToken: string;
+  }): Promise<RawLoginResponseBody | null> {
+    return this.authenticate("authenticateWithAccessToken", args);
+  }
+
+  private authenticate(
+    method: "authenticateWithAuthCode" | "authenticateWithAccessToken",
+    args: string | { accessToken: string; refreshToken: string }
   ): Promise<RawLoginResponseBody | null> {
     this.socket = new WebSocket(this.proxyUrl, this.options);
     this.state = ChannelState.CONNECTING;
@@ -103,7 +118,7 @@ export default class WsJsonClientProxy implements WsJsonClient {
       socket.onopen = () => {
         logger("proxy ws connection opened");
         this.state = ChannelState.CONNECTED;
-        this.doAuthenticate(accessToken).then((res) => {
+        this.doAuthenticate(method, args).then((res) => {
           logger("proxy ws authentication response: %O", res);
           if (isString(res) && res.includes("NOT_AUTHORIZED")) {
             reject(res);
@@ -126,12 +141,10 @@ export default class WsJsonClientProxy implements WsJsonClient {
   }
 
   private doAuthenticate(
-    accessToken: string
+    method: "authenticateWithAuthCode" | "authenticateWithAccessToken",
+    args: string | { accessToken: string; refreshToken: string }
   ): Promise<RawLoginResponseBody | null> {
-    return this.dispatch<RawLoginResponseBody | null>(
-      "authenticate",
-      accessToken
-    ).promise();
+    return this.dispatch<RawLoginResponseBody | null>(method, args).promise();
   }
 
   accountPositions(accountNumber: string): AsyncIterable<PositionsResponse> {