better connection counting

Dockers/xtream-proxy/server.js

@@ -6,7 +6,7 @@ const {
   probeCycle,
   startProbing,
   stopProbing,
-} = require('./src/upstreams');
+} = require('./src/upstream-pool');
 const { handleProxy } = require('./src/proxy');
 const { healthPayload } = require('./src/health');
 const { sendText, sendJson } = require('./src/http-utils');
@@ -16,7 +16,8 @@ const HEALTH_PATHS = new Set(['/health', '/healthz']);
 const server = http.createServer((req, res) => {
   let reqUrl;
   try {
-    reqUrl = new URL(req.url, 'http://proxy.invalid');
+    const rawUrl = req.url.replace(/^\/+/, '/');
+    reqUrl = new URL(rawUrl, 'http://proxy.invalid');
   } catch {
     sendText(res, 400, 'Bad request');
     return;

Dockers/xtream-proxy/src/config.js

@@ -77,7 +77,7 @@ function readConfigEntries() {
     }
 
     for (const url of validUrls) {
-      entries.push({ url, username, password });
+      entries.push({ url, username, password, entryId: index });
     }
   });
   return entries;
@@ -100,7 +100,14 @@ function reloadConfig() {
     const upstream = new Upstream(entry);
     if (seen.has(upstream.key)) continue;
     seen.add(upstream.key);
-    next.push(previous.get(upstream.key) || upstream);
+
+    const existing = previous.get(upstream.key);
+    if (existing) {
+      existing.entryId = upstream.entryId;
+      next.push(existing);
+    } else {
+      next.push(upstream);
+    }
   }
 
   const removed = upstreams.filter((u) => !seen.has(u.key)).length;

Dockers/xtream-proxy/src/health.js

@@ -6,7 +6,8 @@ const {
   getLastProbeAt,
   rankedUpstreams,
   describeUpstream,
-} = require('./upstreams');
+  totalConnectionCapacity,
+} = require('./upstream-pool');
 
 function healthPayload() {
   const upstreams = getUpstreams();
@@ -18,6 +19,7 @@ function healthPayload() {
     configFile: settings.configFile,
     probeIntervalSec: settings.probeIntervalSec,
     lastProbeAt: getLastProbeAt(),
+    maxConnections: totalConnectionCapacity(),
     selected: selected ? describeUpstream(selected.upstream) : null,
     upstreams: upstreams.map(describeUpstream),
   };

Dockers/xtream-proxy/src/proxy.js

@@ -8,7 +8,8 @@ const { buildUpstreamUrl } = require('./upstream');
 const {
   selectUpstream,
   requestProbeSoon,
-} = require('./upstreams');
+  totalConnectionCapacity,
+} = require('./upstream-pool');
 const {
   ReplaceStream,
   buildReplacements,
@@ -39,6 +40,68 @@ function markUpstreamFailure(upstream, error) {
   requestProbeSoon();
 }
 
+const AUTH_RESPONSE_MAX_BYTES = 4 * 1024 * 1024;
+
+function isAuthRequest(req, reqUrl) {
+  if (req.method === 'HEAD') return false;
+  if (reqUrl.searchParams.get('action')) return false;
+  return /\/player_api\.php\/?$/i.test(reqUrl.pathname);
+}
+
+function applyConnectionCapacity(body, upstream, proxyOrigin) {
+  let text = body.toString('utf8');
+
+  try {
+    const json = JSON.parse(text);
+    const info = json && (json.user_info || json.userInfo);
+    if (info && typeof info === 'object') {
+      const max = totalConnectionCapacity();
+      if (max > 0) {
+        info.max_connections =
+          typeof info.max_connections === 'number' ? max : String(max);
+        text = JSON.stringify(json);
+      }
+    }
+  } catch {}
+
+  return Buffer.from(rewriteUrlValue(text, upstream, proxyOrigin));
+}
+
+function sendAuthResponse(proxyRes, res, upstream, proxyOrigin) {
+  const chunks = [];
+  let size = 0;
+  let overflow = false;
+
+  proxyRes.on('data', (chunk) => {
+    if (overflow) {
+      res.write(chunk);
+      return;
+    }
+
+    size += chunk.length;
+    if (size > AUTH_RESPONSE_MAX_BYTES) {
+      overflow = true;
+      for (const buffered of chunks) res.write(buffered);
+      chunks.length = 0;
+      res.write(chunk);
+      return;
+    }
+
+    chunks.push(chunk);
+  });
+
+  proxyRes.on('end', () => {
+    if (overflow) {
+      res.end();
+      return;
+    }
+
+    res.end(
+      applyConnectionCapacity(Buffer.concat(chunks), upstream, proxyOrigin)
+    );
+  });
+}
+
 function handleProxy(req, res, reqUrl) {
   const creds = extractCreds(reqUrl);
 
@@ -182,6 +245,9 @@ function handleProxy(req, res, reqUrl) {
           });
 
           if (rewriting) {
+            if (isAuthRequest(req, reqUrl)) {
+              sendAuthResponse(proxyRes, res, upstream, proxyOrigin);
+            } else {
               const replacer = new ReplaceStream(
                 buildReplacements(upstream, proxyOrigin)
               );
@@ -192,6 +258,7 @@ function handleProxy(req, res, reqUrl) {
                 release();
               });
               proxyRes.pipe(replacer).pipe(res);
+            }
           } else {
             proxyRes.pipe(res);
           }

Dockers/xtream-proxy/src/upstream-pool.js

@@ -25,6 +25,34 @@ function isBusy(upstream) {
   return effectiveLoad(upstream) >= max;
 }
 
+function totalConnectionCapacity() {
+  const perEntry = new Map();
+
+  for (const upstream of getUpstreams()) {
+    const entryId = upstream.entryId ?? upstream.key;
+    const current = perEntry.get(entryId);
+    const maxCons = upstream.state.maxCons;
+
+    if (!current) {
+      perEntry.set(entryId, { maxCons });
+      continue;
+    }
+
+    if (maxCons !== null && maxCons > 0) {
+      current.maxCons =
+        current.maxCons === null || current.maxCons <= 0
+          ? maxCons
+          : Math.max(current.maxCons, maxCons);
+    }
+  }
+
+  let total = 0;
+  for (const { maxCons } of perEntry.values()) {
+    total += maxCons !== null && maxCons > 0 ? maxCons : 1;
+  }
+  return total;
+}
+
 function rankedUpstreams() {
   return getUpstreams()
     .filter((u) => u.state.healthy)
@@ -158,6 +186,7 @@ module.exports = {
   selectUpstream,
   rankedUpstreams,
   isBusy,
+  totalConnectionCapacity,
   describeUpstream,
   getLastProbeAt,
 };

Dockers/xtream-proxy/src/upstream.js

@@ -19,7 +19,7 @@ function baseState() {
 }
 
 class Upstream {
-  constructor({ url, username, password }) {
+  constructor({ url, username, password, entryId }) {
     const parsed = new URL(url);
     parsed.hash = '';
     parsed.search = '';
@@ -27,6 +27,7 @@ class Upstream {
     this.basePath = parsed.pathname.replace(/\/+$/, '');
     this.username = username;
     this.password = password;
+    this.entryId = entryId;
     this.key = `${parsed.origin}${this.basePath}|${username}|${password}`;
     this.state = { ...baseState(), inFlight: 0 };
   }
@@ -41,9 +42,10 @@ class Upstream {
 }
 
 function buildUpstreamUrl(upstream, pathname, search) {
-  const p = pathname.startsWith('/') ? pathname : `/${pathname}`;
+  const basePath = upstream.basePath.replace(/\/+/g, '/');
+  const suffix = `/${pathname}`.replace(/\/+/g, '/');
   const target = new URL(upstream.base.origin);
-  target.pathname = `${upstream.basePath}${p}` || '/';
+  target.pathname = `${basePath}${suffix}` || '/';
   target.search = search || '';
   return target;
 }