Fix sig extraction in main variant of player 4e51e895 (#48)

src/yt/solver/sig.ts

@@ -130,103 +130,128 @@ const identifier: DeepPartial<ESTree.Node> = {
 export function extract(
   node: ESTree.Node,
 ): ESTree.ArrowFunctionExpression | null {
-  if (!matchesStructure(node, identifier)) {
-    return null;
-  }
-  let block: ESTree.BlockStatement | undefined | null;
-  if (
+  const blocks: ESTree.BlockStatement[] = [];
+
+  if (matchesStructure(node, identifier)) {
+    if (
+      node.type === "ExpressionStatement" &&
+      node.expression.type === "AssignmentExpression" &&
+      node.expression.right.type === "FunctionExpression" &&
+      node.expression.right.params.length === 3
+    ) {
+      blocks.push(node.expression.right.body!);
+    } else if (node.type === "VariableDeclaration") {
+      for (const decl of node.declarations) {
+        if (
+          decl.init?.type === "FunctionExpression" &&
+          decl.init.params.length === 3
+        ) {
+          blocks.push(decl.init.body!);
+        }
+      }
+    } else if (
+      node.type === "FunctionDeclaration" &&
+      node.params.length === 3
+    ) {
+      blocks.push(node.body!);
+    } else {
+      return null;
+    }
+  } else if (
     node.type === "ExpressionStatement" &&
-    node.expression.type === "AssignmentExpression" &&
-    node.expression.right.type === "FunctionExpression"
+    node.expression.type === "SequenceExpression"
   ) {
-    block = node.expression.right.body;
-  } else if (node.type === "VariableDeclaration") {
-    for (const decl of node.declarations) {
+    for (const expr of node.expression.expressions) {
       if (
-        decl.type === "VariableDeclarator" &&
-        decl.init?.type === "FunctionExpression" &&
-        decl.init?.params.length === 3
+        expr.type === "AssignmentExpression" &&
+        expr.right.type === "FunctionExpression" &&
+        expr.right.params.length === 3
       ) {
-        block = decl.init.body;
-        break;
+        blocks.push(expr.right.body as ESTree.BlockStatement);
       }
     }
-  } else if (node.type === "FunctionDeclaration") {
-    block = node.body;
   } else {
     return null;
   }
-  const relevantExpression = block?.body.at(-2);
 
-  let call: ESTree.CallExpression | null = null;
-  if (matchesStructure(relevantExpression!, logicalExpression)) {
-    if (
-      relevantExpression?.type !== "ExpressionStatement" ||
-      relevantExpression.expression.type !== "LogicalExpression" ||
-      relevantExpression.expression.right.type !== "SequenceExpression" ||
-      relevantExpression.expression.right.expressions[0].type !==
-        "AssignmentExpression" ||
-      relevantExpression.expression.right.expressions[0].right.type !==
-        "CallExpression"
-    ) {
-      return null;
-    }
-    call = relevantExpression.expression.right.expressions[0].right;
-  } else if (
-    relevantExpression?.type === "IfStatement" &&
-    relevantExpression.consequent.type === "BlockStatement"
-  ) {
-    for (const n of relevantExpression.consequent.body) {
-      if (!matchesStructure(n, nsigExpression)) {
-        continue;
+  for (const block of blocks) {
+    let call: ESTree.CallExpression | null = null;
+
+    for (const stmt of block.body) {
+      if (matchesStructure(stmt, logicalExpression)) {
+        if (
+          stmt.type === "ExpressionStatement" &&
+          stmt.expression.type === "LogicalExpression" &&
+          stmt.expression.right.type === "SequenceExpression" &&
+          stmt.expression.right.expressions[0].type ===
+            "AssignmentExpression" &&
+          stmt.expression.right.expressions[0].right.type === "CallExpression"
+        ) {
+          call = stmt.expression.right.expressions[0].right;
+          break;
+        }
+      } else if (stmt.type === "IfStatement") {
+        let consequent = stmt.consequent;
+        while (consequent.type === "LabeledStatement") {
+          consequent = consequent.body;
+        }
+
+        if (consequent.type === "BlockStatement") {
+          for (const n of consequent.body) {
+            if (!matchesStructure(n, nsigExpression)) {
+              continue;
+            }
+
+            if (
+              n.type === "VariableDeclaration" &&
+              n.declarations[0]?.init?.type === "CallExpression"
+            ) {
+              call = n.declarations[0].init;
+              break;
+            }
+          }
+        }
+        if (call) break;
       }
-      if (
-        n.type !== "VariableDeclaration" ||
-        n.declarations[0].init?.type !== "CallExpression"
-      ) {
-        continue;
-      }
-      call = n.declarations[0].init;
-      break;
     }
-  }
-  if (call === null) {
-    return null;
-  }
-  // TODO: verify identifiers here
-  return {
-    type: "ArrowFunctionExpression",
-    params: [
-      {
-        type: "Identifier",
-        name: "sig",
+
+    if (call?.callee.type !== "Identifier") {
+      continue;
+    }
+
+    // TODO: verify identifiers here
+    return {
+      type: "ArrowFunctionExpression",
+      params: [
+        {
+          type: "Identifier",
+          name: "sig",
+        },
+      ],
+      body: {
+        type: "CallExpression",
+        callee: {
+          type: "Identifier",
+          name: call.callee.name,
+        },
+        arguments: call.arguments.map((arg): ESTree.Expression => {
+          if (
+            arg.type === "CallExpression" &&
+            arg.callee.type === "Identifier" &&
+            arg.callee.name === "decodeURIComponent" &&
+            arg.arguments[0]?.type === "Identifier"
+          ) {
+            return { type: "Identifier", name: "sig" };
+          }
+          return arg as unknown as ESTree.Expression;
+        }),
+        optional: false,
       },
-    ],
-    body: {
-      type: "CallExpression",
-      callee: {
-        type: "Identifier",
-        name: call.callee.name,
-      },
-      arguments:
-        call.arguments.length === 1
-          ? [
-              {
-                type: "Identifier",
-                name: "sig",
-              },
-            ]
-          : [
-              call.arguments[0],
-              {
-                type: "Identifier",
-                name: "sig",
-              },
-            ],
-      optional: false,
-    },
-    async: false,
-    expression: false,
-    generator: false,
-  };
+      async: false,
+      expression: false,
+      generator: false,
+    };
+  }
+
+  return null;
 }

src/yt/solver/test/tests.ts

@@ -286,6 +286,8 @@ export const tests: {
   {
     // tce variant broke sig solving; n and other variants are added only for regression testing
     player: "c1c87fb0",
+    // TODO: fix other variants in other PR
+    variants: ["main", "tcc", "es5", "es6", "tv", "tv_es6", "phone"],
     n: [
       // Synthetic test
       { input: "ZdZIqFPQK-Ty8wId", expected: "jCHBK5GuAFNa2" },
@@ -300,6 +302,23 @@ export const tests: {
       },
     ],
   },
+  {
+    player: "4e51e895",
+    variants: ["main"],
+    n: [
+      // Synthetic test
+      { input: "0eRGgQWJGfT5rFHFj", expected: "t5kO23_msekBur" },
+    ],
+    sig: [
+      {
+        // Synthetic test
+        input:
+          "AL6p_8AwdY9yAhRzK8rYA_9n97Kizf7_9n97Kizf7_9n97Kizf7_9n97Kizf7_9n97Kizf7_9n97Kizf7",
+        expected:
+          "AwdY9yAhRzK8rYA_9n97Kizf7_9n97Kizf7_9n9pKizf7_9n97Kizf7_9n97Kizf7_9n97Kizf7",
+      },
+    ],
+  },
 ];
 
 export const players = new Map([