1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187

const MAX_RETRIES = 3;
const RETRY_DELAY_MS = 500;
const EXPIRATION_BUFFER_SECONDS = 60;

async function getJwtToken(HF_SPACE_NAME, HF_TOKEN, HF_SPACE_USER, D1, forceRefresh = false) {
  const now = Date.now() / 1000;
  let isRefreshing;
  let retries = 0;
  let jwtToken = null;
  let jwtExpiration = 0;

  while (true) {
    let result;
    try {
      result = await D1.prepare("SELECT token, expiration, is_refreshing FROM tokens WHERE id = 1").first();
    } catch (e) {
      throw new Error("Database access error");
    }

    if (result) {
      jwtToken = result.token;
      jwtExpiration = result.expiration;
      isRefreshing = result.is_refreshing;
    } else {
      jwtToken = null;
      jwtExpiration = 0;
      isRefreshing = 0;
    }

    if (jwtToken && jwtExpiration > now + EXPIRATION_BUFFER_SECONDS && !forceRefresh) {
      return jwtToken;
    }

    if (isRefreshing) {
      await new Promise((resolve) => setTimeout(resolve, 100));
      continue;
    }

    try {
      await D1.prepare("INSERT OR REPLACE INTO tokens (id, is_refreshing) VALUES (1, 1)").run();
    } catch (e) {
      throw new Error("Database update error");
    }

    try {
      if (!HF_TOKEN || !HF_SPACE_NAME || !HF_SPACE_USER) {
        throw new Error('One or more required environment variables are missing.');
      }

      const HF_API_URL = `https://huggingface.co/api/spaces/${HF_SPACE_USER}/${HF_SPACE_NAME}/jwt`;

      let response;
      while (retries < MAX_RETRIES) {
        try {
          response = await fetch(HF_API_URL, {
            headers: { "Authorization": `Bearer ${HF_TOKEN}` },
          });

          if (!response.ok) {
            if (response.status >= 500 && response.status < 600) {
              retries++;
              await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY_MS));
              continue;
            } else {
              const errorText = await response.text();
              throw new Error(`Failed to fetch JWT token: ${response.status} ${response.statusText} - ${errorText}`);
            }
          }

          break;

        } catch (networkError) {
          retries++;
          await new Promise((resolve) => setTimeout(resolve, RETRY_DELAY_MS));
        }
      }

      if (retries === MAX_RETRIES) {
        throw new Error("Max retries reached while fetching JWT token");
      }

      const apiResult = await response.json();
      jwtToken = apiResult.token;

      try {
        const jwtPayload = JSON.parse(atob(jwtToken.split('.')[1]));
        jwtExpiration = jwtPayload.exp;
      } catch (_) {
        jwtExpiration = now + 3600;
      }

      try {
        await D1.prepare("INSERT OR REPLACE INTO tokens (id, token, expiration, is_refreshing) VALUES (1, ?, ?, 0)")
          .bind(jwtToken, jwtExpiration)
          .run();
      } catch (e) {
        throw new Error("Database update error");
      }

      return jwtToken;

    } catch (e) {
      try {
        await D1.prepare("UPDATE tokens SET is_refreshing = 0 WHERE id = 1").run();
      } catch (dbError) {}
      throw e;
    }
  }
}

async function initDatabase(D1) {
  try {
    await D1.batch([
      D1.prepare(`CREATE TABLE IF NOT EXISTS tokens (
        id INTEGER PRIMARY KEY,
        token TEXT,
        expiration REAL,
        is_refreshing INTEGER DEFAULT 0
      );`).bind()
    ]);
  } catch (e) {
    throw new Error("Failed to initialize database");
  }
}

export default {
  async fetch(request, env) {
    const HF_TOKEN = env.HF_TOKEN;
    const HF_SPACE_NAME = env.HF_SPACE_NAME;
    const HF_SPACE_USER = env.HF_SPACE_USER;
    const D1 = env.D1;

    try {
      await initDatabase(D1);
      let token = await getJwtToken(HF_SPACE_NAME, HF_TOKEN, HF_SPACE_USER, D1);

      const url = new URL(request.url);
      url.host = `${HF_SPACE_USER}-${HF_SPACE_NAME}.hf.space`;

      const headers = new Headers();

      for (const [key, value] of request.headers.entries()) {
        headers.set(key, value);
      }

      const originalCookies = headers.get('Cookie') || headers.get('cookie') || '';
      const cookieString = originalCookies ? `${originalCookies}; spaces-jwt=${token}` : `spaces-jwt=${token}`;
      headers.set('Cookie', cookieString);

      let newRequest = new Request(url.toString(), {
        method: request.method,
        headers: headers,
        body: request.body,
        redirect: request.redirect,
      });

      let response = await fetch(newRequest);

      if (response.status === 401 || response.status === 403) {
        token = await getJwtToken(HF_SPACE_NAME, HF_TOKEN, HF_SPACE_USER, D1, true);
        const updatedCookieString = originalCookies ? `${originalCookies}; spaces-jwt=${token}` : `spaces-jwt=${token}`;
        headers.set('Cookie', updatedCookieString);

        newRequest = new Request(url.toString(), {
          method: request.method,
          headers: headers,
          body: request.body,
          redirect: request.redirect,
        });

        response = await fetch(newRequest);
      }

      const modifiedHeaders = new Headers(response.headers);
      modifiedHeaders.delete('Link');

      return new Response(response.body, {
        status: response.status,
        statusText: response.statusText,
        headers: modifiedHeaders,
      });

    } catch (error) {
      return new Response(`Error: ${error.message}`, { status: 500 });
    }
  },
};