1
// Copyright (c) 2023 Cloudflare, Inc.
2
// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3
//     https://opensource.org/licenses/Apache-2.0
4
 
5
// A simple test to confirm we can close() a websocket from the close handler.
6
export class DurableObjectExample {
7
  constructor(state) {
8
    this.reset = true;
9
    this.state = state;
10
  }
11
 
12
  async fetch(request) {
13
    // Confirm this is a websocket request.
14
    const upgradeHeader = request.headers.get('Upgrade');
15
    if (!upgradeHeader || upgradeHeader !== 'websocket') {
16
      return new Response('Expected Upgrade: websocket', { status: 426 });
17
    }
18
 
19
    let pair = new WebSocketPair();
20
    let server = pair[0];
21
    if (request.url.endsWith('/hibernation')) {
22
      this.state.acceptWebSocket(server);
23
    } else {
24
      server.accept();
25
      server.addEventListener('message', () => {
26
        server.send('regular message from DO');
27
      });
28
      server.addEventListener('close', () => {
29
        server.close(1000, 'regular close from DO');
30
      });
31
    }
32
 
33
    return new Response(null, {
34
      status: 101,
35
      webSocket: pair[1],
36
    });
37
  }
38
 
39
  webSocketMessage(ws) {
40
    ws.send(`Hibernatable message from DO.`);
41
  }
42
 
43
  webSocketClose(ws, code, reason, wasClean) {
44
    ws.close(1000, 'Hibernatable close from DO');
45
  }
46
}
47
 
48
export const test = {
49
  async test(ctrl, env, ctx) {
50
    let id = env.ns.idFromName('foo');
51
    let obj = env.ns.get(id);
52
 
53
    // Test to make sure that we can call ws.close() from the DO's close handler.
54
    let webSocketTest = async (obj, url, expected) => {
55
      let req = await obj.fetch(url, {
56
        headers: {
57
          Upgrade: 'websocket',
58
        },
59
      });
60
 
61
      let ws = req.webSocket;
62
      if (!ws) {
63
        return new Error('Failed to get ws');
64
      }
65
      ws.accept();
66
      let prom = new Promise((resolve, reject) => {
67
        ws.addEventListener('close', (close) => {
68
          if ((close.code != 1000) & (close.reason != expected)) {
69
            reject(`got ${close.reason}`);
70
          }
71
          resolve();
72
        });
73
      });
74
 
75
      ws.send('Hi from Worker!');
76
      ws.close(1000, 'bye from Worker!');
77
 
78
      await prom;
79
    };
80
 
81
    // Normal websocket
82
    await webSocketTest(obj, 'http://example.com/', 'regular close from DO');
83
    // Hibernatable Websocket.
84
    await webSocketTest(
85
      obj,
86
      'http://example.com/hibernation',
87
      'Hibernatable close from DO'
88
    );
89
  },
90
};