Code Coverage
 
Lines
Functions and Methods
Classes and Traits
Total
100.00%
46 / 46
100.00%
5 / 5
 CRAP
100.00%
1 / 1
REST_Jetpack_AI_JWT
100.00%
46 / 46
100.00%
5 / 5
 9
100.00%
1 / 1
 __construct
100.00%
2 / 2
100.00%
1 / 1
 1
 register_rest_route
100.00%
13 / 13
100.00%
1 / 1
 2
 is_rest_endpoint_registered
100.00%
4 / 4
100.00%
1 / 1
 1
 permission_callback
100.00%
1 / 1
100.00%
1 / 1
 2
 get_jwt
100.00%
26 / 26
100.00%
1 / 1
 3
1<?php
2/**
3 * REST_Jetpack_AI_JWT file.
4 *
5 * @package automattic/jetpack-connection
6 */
7
8namespace Automattic\Jetpack\Connection;
9
10use Jetpack_Options;
11
12/**
13 * Registers the `jetpack/v4/jetpack-ai-jwt` route, which asks WordPress.com for a JWT that
14 * Jetpack AI clients use to call the AI completion service.
15 *
16 * @since 9.1.0
17 */
18class REST_Jetpack_AI_JWT extends \WP_REST_Controller {
19
20    /**
21     * REST_Jetpack_AI_JWT constructor.
22     */
23    public function __construct() {
24        $this->namespace = 'jetpack/v4';
25        $this->rest_base = '/jetpack-ai-jwt';
26    }
27
28    /**
29     * Register available routes.
30     */
31    public function register_rest_route() {
32        /*
33         * Check if the `jetpack/v4/jetpack-ai-jwt` endpoint is registered
34         * by the Jetpack plugin to avoid registering it again.
35         * In case it's not registered, register it
36         * to make it available for Jetpack products that depend on it.
37         */
38        if ( $this->is_rest_endpoint_registered() ) {
39            return;
40        }
41
42        register_rest_route(
43            $this->namespace,
44            $this->rest_base,
45            array(
46                array(
47                    'methods'             => \WP_REST_Server::EDITABLE,
48                    'callback'            => array( $this, 'get_jwt' ),
49                    'permission_callback' => array( $this, 'permission_callback' ),
50                ),
51            )
52        );
53    }
54
55    /**
56     * Check if this REST endpoint is already registered.
57     *
58     * @return bool True if the endpoint is registered, false otherwise.
59     */
60    private function is_rest_endpoint_registered() {
61        $server        = rest_get_server();
62        $routes        = $server->get_routes();
63        $full_endpoint = '/' . trim( $this->namespace, '/' ) . $this->rest_base;
64
65        return isset( $routes[ $full_endpoint ] );
66    }
67
68    /**
69     * Permission callback for the JWT endpoint.
70     *
71     * @return bool
72     */
73    public function permission_callback() {
74        return ( new Manager( 'jetpack' ) )->is_user_connected() && current_user_can( 'edit_posts' );
75    }
76
77    /**
78     * Ask WPCOM for a JWT token to use for OpenAI completion.
79     *
80     * @return \WP_REST_Response|\WP_Error The token and blog ID, or the error from WPCOM.
81     */
82    public function get_jwt() {
83        $blog_id = Jetpack_Options::get_option( 'id' );
84
85        $response = Client::wpcom_json_api_request_as_user(
86            "/sites/$blog_id/jetpack-openai-query/jwt",
87            '2',
88            array(
89                'method'  => 'POST',
90                'headers' => array( 'Content-Type' => 'application/json; charset=utf-8' ),
91            ),
92            wp_json_encode( array(), JSON_UNESCAPED_SLASHES | JSON_UNESCAPED_UNICODE ),
93            'wpcom'
94        );
95
96        if ( is_wp_error( $response ) ) {
97            return $response;
98        }
99
100        $json = json_decode( wp_remote_retrieve_body( $response ) );
101
102        if ( ! isset( $json->token ) ) {
103            return new \WP_Error(
104                'no-token',
105                'No token returned from WPCOM',
106                array( 'status' => 500 )
107            );
108        }
109
110        return rest_ensure_response(
111            array(
112                'token'   => $json->token,
113                'blog_id' => $blog_id,
114            )
115        );
116    }
117}