Skip to content

Commit 96628d4

Browse files
Alex-MonahanAlex-Monahan
committed
Store primary key in secret, generate token each request
1 parent ee5b4ea commit 96628d4

File tree

8 files changed

+9661
-10
lines changed

8 files changed

+9661
-10
lines changed

‎CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ set(EXTENSION_SOURCES
2121
src/gsheets_requests.cpp
2222
src/gsheets_read.cpp
2323
src/gsheets_utils.cpp
24+
src/gsheets_get_token.cpp
2425
)
2526

2627
build_static_extension(${TARGET_NAME} ${EXTENSION_SOURCES})

‎src/gsheets_auth.cpp‎

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,9 @@
66
#include "duckdb/main/extension_util.hpp"
77
#include <fstream>
88
#include <cstdlib>
9+
#include <json.hpp>
10+
11+
using json = nlohmann::json;
912

1013
namespace duckdb
1114
{
@@ -78,6 +81,23 @@ namespace duckdb
7881
return std::move(result);
7982
}
8083

84+
// TODO: Maybe this should be a KeyValueSecret
85+
static unique_ptr<BaseSecret> CreateGsheetSecretFromPrivateKey(ClientContext &context, CreateSecretInput &input) {
86+
auto scope = input.scope;
87+
88+
auto result = make_uniq<KeyValueSecret>(scope, input.type, input.provider, input.name);
89+
90+
// Manage specific secret option
91+
CopySecret("secret", input, *result);
92+
CopySecret("email", input, *result);
93+
94+
// Redact sensible keys
95+
RedactCommonKeys(*result);
96+
result->redact_keys.insert("secret");
97+
98+
return std::move(result);
99+
}
100+
81101
void CreateGsheetSecretFunctions::Register(DatabaseInstance &instance)
82102
{
83103
string type = "gsheet";
@@ -100,6 +120,13 @@ namespace duckdb
100120
oauth_function.named_parameters["use_oauth"] = LogicalType::BOOLEAN;
101121
RegisterCommonSecretParameters(oauth_function);
102122
ExtensionUtil::RegisterFunction(instance, oauth_function);
123+
124+
// Register the private key secret provider
125+
CreateSecretFunction private_key_function = {type, "private_key", CreateGsheetSecretFromPrivateKey};
126+
private_key_function.named_parameters["email"] = LogicalType::VARCHAR;
127+
private_key_function.named_parameters["secret"] = LogicalType::VARCHAR;
128+
RegisterCommonSecretParameters(private_key_function);
129+
ExtensionUtil::RegisterFunction(instance, private_key_function);
103130
}
104131

105132
std::string InitiateOAuthFlow()

‎src/gsheets_copy.cpp‎

Lines changed: 35 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@
77
#include "duckdb/common/file_system.hpp"
88
#include "duckdb/main/secret/secret_manager.hpp"
99
#include <json.hpp>
10+
#include <regex>
11+
#include "gsheets_get_token.hpp"
12+
#include <iostream>
1013

1114
using json = nlohmann::json;
1215

@@ -49,12 +52,39 @@ namespace duckdb
4952
throw InvalidInputException("Invalid secret format for 'gsheet' secret");
5053
}
5154

52-
Value token_value;
53-
if (!kv_secret->TryGetValue("token", token_value)) {
54-
throw InvalidInputException("'token' not found in 'gsheet' secret");
55-
}
55+
std::string token;
56+
57+
if (secret.GetProvider() == "private_key") {
58+
// If using a private key, retrieve the private key from the secret, but convert it
59+
// into a token before use. This is an extra request per Google Sheet read or copy.
60+
Value email_value;
61+
if (!kv_secret->TryGetValue("email", email_value)) {
62+
throw InvalidInputException("'email' not found in 'gsheet' secret");
63+
}
64+
std::string email_string = email_value.ToString();
5665

57-
std::string token = token_value.ToString();
66+
Value secret_value;
67+
if (!kv_secret->TryGetValue("secret", secret_value)) {
68+
throw InvalidInputException("'secret' not found in 'gsheet' secret");
69+
}
70+
std::string secret_string = std::regex_replace(secret_value.ToString(), std::regex(R"(\\n)"), "\n");
71+
72+
json token_json = parseJson(get_token(email_string, secret_string));
73+
74+
json failure_string = "failure";
75+
if (token_json["status"] == failure_string) {
76+
throw InvalidInputException("Conversion from private key to token failed. Check key format (-----BEGIN PRIVATE KEY-----\\n ... -----END PRIVATE KEY-----\\n) and expiration date.");
77+
}
78+
79+
token = token_json["access_token"].get<std::string>();
80+
} else {
81+
Value token_value;
82+
if (!kv_secret->TryGetValue("token", token_value)) {
83+
throw InvalidInputException("'token' not found in 'gsheet' secret");
84+
}
85+
86+
token = token_value.ToString();
87+
}
5888
std::string spreadsheet_id = extract_spreadsheet_id(file_path);
5989
std::string sheet_id = extract_sheet_id(file_path);
6090
std::string sheet_name = "Sheet1";

‎src/gsheets_get_token.cpp‎

Lines changed: 142 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,142 @@
1+
// Taken with modifications from https://gist.github.com/niuk/6365b819a86a7e0b92d82328fcf87da5
2+
#include <stdio.h>
3+
#include <string.h>
4+
#include <stdlib.h>
5+
6+
#include <openssl/ssl.h>
7+
#include <openssl/err.h>
8+
9+
#define CPPHTTPLIB_OPENSSL_SUPPORT
10+
#include <httplib.hpp>
11+
12+
#include <json.hpp>
13+
using json = nlohmann::json;
14+
namespace duckdb
15+
{
16+
17+
char get_base64_char(char byte) {
18+
if (byte < 26) {
19+
return 'A' + byte;
20+
} else if (byte < 52) {
21+
return 'a' + byte - 26;
22+
} else if (byte < 62) {
23+
return '0' + byte - 52;
24+
} else if (byte == 62) {
25+
return '-';
26+
} else if (byte == 63) {
27+
return '_';
28+
} else {
29+
fprintf(stderr, "BAD BYTE: %02x\n", byte);
30+
exit(1);
31+
return 0;
32+
}
33+
}
34+
35+
// To execute C, please define "int main()"
36+
void base64encode(char *output, const char *input, size_t input_length) {
37+
size_t input_index = 0;
38+
size_t output_index = 0;
39+
for (; input_index < input_length; ++output_index) {
40+
switch (output_index % 4) {
41+
case 0:
42+
output[output_index] = get_base64_char((0xfc & input[input_index]) >> 2);
43+
break;
44+
case 1:
45+
output[output_index] = get_base64_char(((0x03 & input[input_index]) << 4) | ((0xf0 & input[input_index + 1]) >> 4));
46+
++input_index;
47+
break;
48+
case 2:
49+
output[output_index] = get_base64_char(((0x0f & input[input_index]) << 2) | ((0xc0 & input[input_index + 1]) >> 6));
50+
++input_index;
51+
break;
52+
case 3:
53+
output[output_index] = get_base64_char(0x3f & input[input_index]);
54+
++input_index;
55+
break;
56+
default:
57+
exit(1);
58+
}
59+
}
60+
61+
output[output_index] = '\0';
62+
}
63+
64+
std::string get_token(const std::string& email, const std::string& private_key) {
65+
const char *header = "{\"alg\":\"RS256\",\"typ\":\"JWT\"}";
66+
67+
const int iat = time(NULL);
68+
const int exp = iat + 60 * 60 - 1;
69+
70+
char claim_set[1024];
71+
72+
/* Create jwt claim set */
73+
json jwt_claim_set;
74+
std::time_t t = std::time(NULL);
75+
jwt_claim_set["iss"] = email; /* service account email address */
76+
jwt_claim_set["scope"] = "https://www.googleapis.com/auth/spreadsheets" /* scope of requested access token */;
77+
jwt_claim_set["aud"] = "https://accounts.google.com/o/oauth2/token"; /* intended target of the assertion for an access token */
78+
jwt_claim_set["iat"] = std::to_string(t); /* issued time */
79+
// Since we get a new token on every request (and max request time is 3 minutes),
80+
// set the limit to 10 minutes. (Max that Google allows is 1 hour)
81+
jwt_claim_set["exp"] = std::to_string(t+600); /* expire time*/
82+
83+
std::copy(jwt_claim_set.dump().cbegin(), jwt_claim_set.dump().cend(), claim_set);
84+
85+
char header_64[1024];
86+
base64encode(header_64, header, strlen(header));
87+
88+
char claim_set_64[1024];
89+
base64encode(claim_set_64, claim_set, strlen(claim_set));
90+
91+
char input[1024];
92+
int input_length = sprintf(input, "%s.%s", header_64, claim_set_64);
93+
94+
unsigned char *digest = SHA256((const unsigned char *)input, input_length, NULL);
95+
char digest_str[1024];
96+
for (int i = 0; i < SHA256_DIGEST_LENGTH; ++i) {
97+
sprintf(digest_str + i * 2, "%02x", digest[i]);
98+
}
99+
100+
digest_str[SHA256_DIGEST_LENGTH * 2] = '\0';
101+
102+
// This works, but requires a file. The in-memory version is below.
103+
// FILE *key_file = fopen("credentials.key.pem", "r");
104+
// RSA *rsa = PEM_read_RSAPrivateKey(key_file, NULL, NULL, NULL);
105+
106+
BIO *bufio;
107+
RSA *rsa;
108+
109+
bufio = BIO_new_mem_buf(private_key.c_str(), -1);
110+
PEM_read_bio_RSAPrivateKey(bufio, &rsa, 0, NULL);
111+
112+
if (rsa != NULL) {
113+
unsigned char sigret[4096];
114+
unsigned int siglen;
115+
if (RSA_sign(NID_sha256, digest, SHA256_DIGEST_LENGTH, sigret, &siglen, rsa)) {
116+
if (RSA_verify(NID_sha256, digest, SHA256_DIGEST_LENGTH, sigret, siglen, rsa)) {
117+
char signature_64[1024];
118+
base64encode(signature_64, (const char *)sigret, siglen);
119+
120+
char jwt[1024];
121+
sprintf(jwt, "%s.%s", input, signature_64);
122+
123+
duckdb_httplib_openssl::Client cli("https://oauth2.googleapis.com");
124+
duckdb_httplib_openssl::Params params;
125+
params.emplace("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer");
126+
params.emplace("assertion", jwt);
127+
auto result = cli.Post("/token", params);
128+
return (result -> body);
129+
} else {
130+
printf("Could not verify RSA signature.");
131+
}
132+
} else {
133+
unsigned long err = ERR_get_error();
134+
printf("RSA_sign failed: %lu, %s\n", err, ERR_error_string(err, NULL));
135+
}
136+
137+
RSA_free(rsa);
138+
}
139+
std::string failure_message = "{\"status\":\"failed\"}";
140+
return failure_message;
141+
}
142+
}

‎src/gsheets_read.cpp‎

Lines changed: 34 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,9 @@
55
#include "duckdb/main/secret/secret_manager.hpp"
66
#include "gsheets_requests.hpp"
77
#include <json.hpp>
8+
#include <regex>
9+
#include "gsheets_get_token.hpp"
10+
#include <iostream>
811

912
namespace duckdb {
1013

@@ -113,13 +116,39 @@ unique_ptr<FunctionData> ReadSheetBind(ClientContext &context, TableFunctionBind
113116
throw InvalidInputException("Invalid secret format for 'gsheet' secret");
114117
}
115118

116-
Value token_value;
117-
if (!kv_secret->TryGetValue("token", token_value)) {
118-
throw InvalidInputException("'token' not found in 'gsheet' secret");
119-
}
119+
std::string token;
120+
121+
if (secret.GetProvider() == "private_key") {
122+
// If using a private key, retrieve the private key from the secret, but convert it
123+
// into a token before use. This is an extra request per Google Sheet read or copy.
124+
Value email_value;
125+
if (!kv_secret->TryGetValue("email", email_value)) {
126+
throw InvalidInputException("'email' not found in 'gsheet' secret");
127+
}
128+
std::string email_string = email_value.ToString();
129+
130+
Value secret_value;
131+
if (!kv_secret->TryGetValue("secret", secret_value)) {
132+
throw InvalidInputException("'secret' not found in 'gsheet' secret");
133+
}
134+
std::string secret_string = std::regex_replace(secret_value.ToString(), std::regex(R"(\\n)"), "\n");
135+
136+
json token_json = parseJson(get_token(email_string, secret_string));
137+
138+
json failure_string = "failure";
139+
if (token_json["status"] == failure_string) {
140+
throw InvalidInputException("Conversion from private key to token failed. Check key format (-----BEGIN PRIVATE KEY-----\\n ... -----END PRIVATE KEY-----\\n) and expiration date.");
141+
}
120142

121-
std::string token = token_value.ToString();
143+
token = token_json["access_token"].get<std::string>();
144+
} else {
145+
Value token_value;
146+
if (!kv_secret->TryGetValue("token", token_value)) {
147+
throw InvalidInputException("'token' not found in 'gsheet' secret");
148+
}
122149

150+
token = token_value.ToString();
151+
}
123152
// Get sheet name from URL
124153
std::string sheet_id = extract_sheet_id(sheet_input);
125154
sheet_name = get_sheet_name_from_id(spreadsheet_id, sheet_id, token);

‎src/include/gsheets_get_token.hpp‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
#pragma once
2+
3+
#include <string>
4+
#include <stdlib.h>
5+
6+
namespace duckdb {
7+
8+
char get_base64_char(char byte);
9+
10+
void base64encode(char *output, const char *input, size_t input_length) ;
11+
12+
std::string get_token(const std::string& email, const std::string& private_key) ;
13+
14+
}

‎third_party/LICENSE‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
The MIT License (MIT)
2+
3+
Copyright (c) 2017 yhirose
4+
5+
Permission is hereby granted, free of charge, to any person obtaining a copy
6+
of this software and associated documentation files (the "Software"), to deal
7+
in the Software without restriction, including without limitation the rights
8+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9+
copies of the Software, and to permit persons to whom the Software is
10+
furnished to do so, subject to the following conditions:
11+
12+
The above copyright notice and this permission notice shall be included in all
13+
copies or substantial portions of the Software.
14+
15+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21+
SOFTWARE.
22+

0 commit comments

Comments
 (0)