Chore: make unifi proxy more generic (#6469)

src/utils/proxy/handlers/unifi.js

@@ -0,0 +1,116 @@
+import cache from "memory-cache";
+
+import createLogger from "utils/logger";
+import { formatApiCall } from "utils/proxy/api-helpers";
+import { addCookieToJar, setCookieHeader } from "utils/proxy/cookie-jar";
+import { httpProxy } from "utils/proxy/http";
+import widgets from "widgets/widgets";
+
+function isSuccessfulLoginResponse(data) {
+  const json = JSON.parse(data.toString());
+  return json?.meta?.rc === "ok" || json?.login_time || json?.update_time;
+}
+
+async function login({ widget, api, endpoint, csrfToken }) {
+  const loginUrl = new URL(formatApiCall(api.replace("{prefix}", ""), { endpoint, ...widget }));
+  const headers = { "Content-Type": "application/json" };
+
+  if (csrfToken) {
+    headers["X-CSRF-TOKEN"] = csrfToken;
+  }
+
+  return httpProxy(loginUrl, {
+    method: "POST",
+    body: JSON.stringify({ username: widget.username, password: widget.password, remember: true, rememberMe: true }),
+    headers,
+  });
+}
+
+export default function createUnifiProxyHandler({
+  proxyName,
+  resolveWidget,
+  resolveRequestContext,
+  getLoginEndpoint = () => "auth/login",
+  shouldAttemptLogin = ({ widget }) => !widget.key,
+}) {
+  const prefixCacheKey = `${proxyName}__prefix`;
+  const logger = createLogger(proxyName);
+
+  return async function unifiProxyHandler(req, res) {
+    const widget = await resolveWidget(req, logger);
+    const { service, endpoint } = req.query;
+
+    if (!widget) {
+      return res.status(400).json({ error: "Invalid proxy service type" });
+    }
+
+    const api = widgets?.[widget.type]?.api;
+    if (!api) {
+      return res.status(403).json({ error: "Service does not support API calls" });
+    }
+
+    const cachedPrefix = cache.get(`${prefixCacheKey}.${service}`);
+    const {
+      prefix,
+      headers = {},
+      csrfToken: initialCsrfToken,
+    } = await resolveRequestContext({
+      cachedPrefix,
+      logger,
+      req,
+      service,
+      widget,
+    });
+    let csrfToken = initialCsrfToken;
+
+    cache.put(`${prefixCacheKey}.${service}`, prefix);
+
+    widget.prefix = prefix;
+    const url = new URL(formatApiCall(api, { endpoint, ...widget }));
+    const params = { method: "GET", headers };
+    setCookieHeader(url, params);
+
+    let [status, contentType, data, responseHeaders] = await httpProxy(url, params);
+
+    if (status === 401 && shouldAttemptLogin({ widget, req, responseHeaders })) {
+      logger.debug("UniFi request was rejected, attempting login.");
+
+      if (responseHeaders?.["x-csrf-token"]) {
+        csrfToken = responseHeaders["x-csrf-token"];
+      }
+
+      [status, contentType, data, responseHeaders] = await login({
+        api,
+        csrfToken,
+        endpoint: getLoginEndpoint({ prefix, req, widget }),
+        widget,
+      });
+
+      if (status !== 200) {
+        logger.error("HTTP %d logging in to UniFi. Data: %s", status, data);
+        return res.status(status).json({ error: { message: `HTTP Error ${status}`, url, data } });
+      }
+
+      if (!isSuccessfulLoginResponse(data)) {
+        logger.error("Error logging in to UniFi: Data: %s", data);
+        return res.status(401).end(data);
+      }
+
+      addCookieToJar(url, responseHeaders);
+      setCookieHeader(url, params);
+
+      [status, contentType, data, responseHeaders] = await httpProxy(url, params);
+    }
+
+    if (status !== 200) {
+      logger.error("HTTP %d getting data from UniFi endpoint %s. Data: %s", status, url.href, data);
+      return res.status(status).json({ error: { message: `HTTP Error ${status}`, url, data } });
+    }
+
+    if (contentType) {
+      res.setHeader("Content-Type", contentType);
+    }
+
+    return res.status(status).send(data);
+  };
+}

src/widgets/unifi/proxy.js

@@ -1,19 +1,11 @@
-import cache from "memory-cache";
-
 import getServiceWidget from "utils/config/service-helpers";
 import { getPrivateWidgetOptions } from "utils/config/widget-helpers";
-import createLogger from "utils/logger";
+import createUnifiProxyHandler from "utils/proxy/handlers/unifi";
-import { formatApiCall } from "utils/proxy/api-helpers";
-import { addCookieToJar, setCookieHeader } from "utils/proxy/cookie-jar";
 import { httpProxy } from "utils/proxy/http";
-import widgets from "widgets/widgets";
 
 const udmpPrefix = "/proxy/network";
-const proxyName = "unifiProxyHandler";
-const prefixCacheKey = `${proxyName}__prefix`;
-const logger = createLogger(proxyName);
 
-async function getWidget(req) {
+async function getWidget(req, logger) {
   const { group, service, index } = req.query;
 
   let widget = null;
@@ -43,100 +35,36 @@ async function getWidget(req) {
   return widget;
 }
 
-async function login(widget, csrfToken) {
+async function resolveRequestContext({ cachedPrefix, widget }) {
-  const endpoint = widget.prefix === udmpPrefix ? "auth/login" : "login";
-  const api = widgets?.[widget.type]?.api?.replace("{prefix}", ""); // no prefix for login url
-  const loginUrl = new URL(formatApiCall(api, { endpoint, ...widget }));
-  const loginBody = { username: widget.username, password: widget.password, remember: true, rememberMe: true };
-  const headers = { "Content-Type": "application/json" };
-  if (csrfToken) {
-    headers["X-CSRF-TOKEN"] = csrfToken;
-  }
-  const [status, contentType, data, responseHeaders] = await httpProxy(loginUrl, {
-    method: "POST",
-    body: JSON.stringify(loginBody),
-    headers,
-  });
-  return [status, contentType, data, responseHeaders];
-}
-
-export default async function unifiProxyHandler(req, res) {
-  const widget = await getWidget(req);
-  const { service } = req.query;
-  if (!widget) {
-    return res.status(400).json({ error: "Invalid proxy service type" });
-  }
-
-  const api = widgets?.[widget.type]?.api;
-  if (!api) {
-    return res.status(403).json({ error: "Service does not support API calls" });
-  }
-
-  let [status, contentType, data, responseHeaders] = [];
-  let prefix = cache.get(`${prefixCacheKey}.${service}`);
-  let csrfToken;
   const headers = {};
+
   if (widget.key) {
-    prefix = udmpPrefix;
     headers["X-API-KEY"] = widget.key;
-    headers["Accept"] = "application/json";
+    headers.Accept = "application/json";
-  } else if (prefix === null) {
+    return { headers, prefix: udmpPrefix };
-    // auto detect if we're talking to a UDM Pro or Network API device, and cache the result
+  }
-    // so that we don't make two requests each time data from Unifi is required
+
-    [status, contentType, data, responseHeaders] = await httpProxy(widget.url);
+  if (cachedPrefix !== null) {
-    prefix = "";
+    return { headers, prefix: cachedPrefix };
+  }
+
+  const [, , , responseHeaders] = await httpProxy(widget.url);
+  let prefix = "";
+  let csrfToken;
+
   if (responseHeaders?.["x-csrf-token"]) {
-      // Unifi OS < 3.2.5 passes & requires csrf-token
     prefix = udmpPrefix;
     csrfToken = responseHeaders["x-csrf-token"];
-    } else if (
+  } else if (responseHeaders?.["access-control-expose-headers"] || responseHeaders?.["Access-Control-Expose-Headers"]) {
-      responseHeaders?.["access-control-expose-headers"] ||
-      responseHeaders?.["Access-Control-Expose-Headers"]
-    ) {
-      // Unifi OS ≥ 3.2.5 doesnt pass csrf token but still uses different endpoint, same with Network API
     prefix = udmpPrefix;
   }
-  }
-  cache.put(`${prefixCacheKey}.${service}`, prefix);
 
-  widget.prefix = prefix;
+  return { csrfToken, headers, prefix };
-  const { endpoint } = req.query;
-  const url = new URL(formatApiCall(api, { endpoint, ...widget }));
-  const params = { method: "GET", headers };
-  setCookieHeader(url, params);
-
-  [status, contentType, data, responseHeaders] = await httpProxy(url, params);
-
-  if (status === 401 && !widget.key) {
-    logger.debug("Unifi isn't logged in or rejected the reqeust, attempting login.");
-    if (responseHeaders?.["x-csrf-token"]) {
-      csrfToken = responseHeaders["x-csrf-token"];
-    }
-    [status, contentType, data, responseHeaders] = await login(widget, csrfToken);
-
-    if (status !== 200) {
-      logger.error("HTTP %d logging in to Unifi. Data: %s", status, data);
-      return res.status(status).json({ error: { message: `HTTP Error ${status}`, url, data } });
-    }
-
-    const json = JSON.parse(data.toString());
-    if (!(json?.meta?.rc === "ok" || json?.login_time || json?.update_time)) {
-      logger.error("Error logging in to Unifi: Data: %s", data);
-      return res.status(401).end(data);
-    }
-
-    addCookieToJar(url, responseHeaders);
-    setCookieHeader(url, params);
-
-    logger.debug("Retrying Unifi request after login.");
-    [status, contentType, data, responseHeaders] = await httpProxy(url, params);
-  }
-
-  if (status !== 200) {
-    logger.error("HTTP %d getting data from Unifi endpoint %s. Data: %s", status, url.href, data);
-    return res.status(status).json({ error: { message: `HTTP Error ${status}`, url, data } });
-  }
-
-  if (contentType) res.setHeader("Content-Type", contentType);
-  return res.status(status).send(data);
 }
+
+export default createUnifiProxyHandler({
+  proxyName: "unifiProxyHandler",
+  resolveWidget: getWidget,
+  resolveRequestContext,
+  getLoginEndpoint: ({ prefix }) => (prefix === udmpPrefix ? "auth/login" : "login"),
+});

src/widgets/unifi/proxy.test.js

@@ -89,4 +89,58 @@ describe("widgets/unifi/proxy", () => {
     expect(res.statusCode).toBe(200);
     expect(res.body).toEqual(Buffer.from("data"));
   });
+
+  it("uses unifi_console private widget config for the info widget path", async () => {
+    getPrivateWidgetOptions.mockResolvedValue({
+      url: "http://console",
+      username: "u",
+      password: "p",
+    });
+
+    httpProxy
+      .mockResolvedValueOnce([200, "text/html", Buffer.from(""), { "x-csrf-token": "csrf" }])
+      .mockResolvedValueOnce([200, "application/json", Buffer.from("data"), {}]);
+
+    const req = {
+      query: {
+        group: "unifi_console",
+        service: "unifi_console",
+        endpoint: "self",
+        query: JSON.stringify({ index: 2 }),
+      },
+    };
+    const res = createMockRes();
+
+    await unifiProxyHandler(req, res);
+
+    expect(getPrivateWidgetOptions).toHaveBeenCalledWith("unifi_console", 2);
+    expect(httpProxy.mock.calls[1][0].toString()).toContain("/proxy/network/api/self");
+    expect(res.statusCode).toBe(200);
+  });
+
+  it("uses the API key flow without attempting login", async () => {
+    getServiceWidget.mockResolvedValue({
+      type: "unifi",
+      key: "secret",
+      url: "http://unifi",
+    });
+
+    httpProxy.mockResolvedValueOnce([200, "application/json", Buffer.from("data"), {}]);
+
+    const req = { query: { group: "g", service: "svc", endpoint: "self", index: "0" } };
+    const res = createMockRes();
+
+    await unifiProxyHandler(req, res);
+
+    expect(httpProxy).toHaveBeenCalledTimes(1);
+    expect(httpProxy.mock.calls[0][0].toString()).toContain("/proxy/network/api/self");
+    expect(httpProxy.mock.calls[0][1]).toMatchObject({
+      headers: {
+        Accept: "application/json",
+        "X-API-KEY": "secret",
+      },
+      method: "GET",
+    });
+    expect(res.statusCode).toBe(200);
+  });
 });