1 #include "avb_tools.h"
2
3 #include "gtest/gtest.h"
4 #include "nugget/app/avb/avb.pb.h"
5
6 #include <app_nugget.h>
7 #include "Avb.client.h"
8 #include <avb.h>
9 #include <nos/NuggetClient.h>
10
11 #include <chrono>
12 #include <iostream>
13 #include <thread>
14 #include <vector>
15
16 #ifdef ANDROID
17 #include <android-base/endian.h>
18 #include "nos/CitadeldProxyClient.h"
19 #else
20 #include "gflags/gflags.h"
21 #endif // ANDROID
22
23 using namespace nugget::app::avb;
24
25 namespace avb_tools {
26
SetBootloader(nos::NuggetClientInterface * client)27 void SetBootloader(nos::NuggetClientInterface *client)
28 {
29 // Force AVB to believe that the AP is in the BIOS.
30 ::nos::AppClient app(*client, APP_ID_AVB_TEST);
31
32 /* We have to have a buffer, because it's called by reference. */
33 std::vector<uint8_t> buffer;
34
35 // No params, no args needed. This is all that the fake "AVB_TEST" app does.
36 uint32_t retval = app.Call(0, buffer, &buffer);
37
38 EXPECT_EQ(retval, APP_SUCCESS);
39 }
40
BootloaderDone(nos::NuggetClientInterface * client)41 void BootloaderDone(nos::NuggetClientInterface *client)
42 {
43 BootloaderDoneRequest request;
44
45 Avb service(*client);
46 ASSERT_NO_ERROR(service.BootloaderDone(request, nullptr), "");
47 }
48
GetState(nos::NuggetClientInterface * client,bool * bootloader,bool * production,uint8_t * locks)49 void GetState(nos::NuggetClientInterface *client, bool *bootloader,
50 bool *production, uint8_t *locks) {
51 GetStateRequest request;
52 GetStateResponse response;
53
54 Avb service(*client);
55 ASSERT_NO_ERROR(service.GetState(request, &response), "");
56 EXPECT_EQ(response.number_of_locks(), 4U);
57
58 if (bootloader != NULL)
59 *bootloader = response.bootloader();
60 if (production != NULL)
61 *production = response.production();
62
63 auto response_locks = response.locks();
64 if (locks != NULL) {
65 for (size_t i = 0; i < response_locks.size(); i++)
66 locks[i] = response_locks[i];
67 }
68 }
69
Reset(nos::NuggetClientInterface * client,ResetRequest_ResetKind kind,const uint8_t * sig,size_t size)70 int Reset(nos::NuggetClientInterface *client, ResetRequest_ResetKind kind,
71 const uint8_t *sig, size_t size) {
72 ResetRequest request;
73
74 request.set_kind(kind);
75 request.mutable_token()->set_selector(ResetToken::CURRENT);
76 if (sig && size) {
77 request.mutable_token()->set_signature(sig, size);
78 } else {
79 uint8_t empty[AVB_SIGNATURE_SIZE];
80 memset(empty, 0, sizeof(empty));
81 request.mutable_token()->set_signature(empty, sizeof(empty));
82 }
83
84 Avb service(*client);
85 return service.Reset(request, nullptr);
86 }
87
GetResetChallenge(nos::NuggetClientInterface * client,uint32_t * selector,uint64_t * nonce,uint8_t * device_data,size_t * len)88 int GetResetChallenge(nos::NuggetClientInterface *client,
89 uint32_t *selector, uint64_t *nonce,
90 uint8_t *device_data, size_t *len) {
91 GetResetChallengeRequest request;
92 GetResetChallengeResponse response;
93
94 Avb service(*client);
95 uint32_t ret = service.GetResetChallenge(request, &response);
96 if (ret != APP_SUCCESS) {
97 return ret;
98 }
99 *selector = response.selector();
100 *nonce = response.nonce();
101 // Only copy what there is space for.
102 *len = (*len < response.device_data().size() ? *len : response.device_data().size());
103 memcpy(device_data, response.device_data().data(), *len);
104 // Let the caller assert if the requested size was too large.
105 *len = response.device_data().size();
106 return ret;
107 }
108
SetProduction(nos::NuggetClientInterface * client,bool production,const uint8_t * data,size_t size)109 int SetProduction(nos::NuggetClientInterface *client, bool production,
110 const uint8_t *data, size_t size) {
111 SetProductionRequest request;
112
113 request.set_production(production);
114 if (size != 0 && data != NULL) {
115 request.set_device_data(data, size);
116 }
117 // Substitute an empty hash
118 uint8_t empty[AVB_DEVICE_DATA_SIZE];
119 memset(empty, 0, sizeof(empty));
120 if (production && data == NULL) {
121 request.set_device_data(empty, sizeof(empty));
122 }
123
124 Avb service(*client);
125 return service.SetProduction(request, nullptr);
126 }
127
ResetProduction(nos::NuggetClientInterface * client)128 void ResetProduction(nos::NuggetClientInterface *client)
129 {
130 struct ResetMessage message;
131 int code;
132 uint32_t selector;
133 size_t len = sizeof(message.data);
134 uint8_t signature[AVB_SIGNATURE_SIZE];
135 size_t siglen = sizeof(signature);
136 memset(signature, 0, siglen);
137
138 // We need the nonce to be set before we get fallthrough.
139 memset(message.data, 0, sizeof(message.data));
140 code = GetResetChallenge(client, &selector, &message.nonce, message.data, &len);
141 ASSERT_NO_ERROR(code, "");
142 // No signature is needed for TEST_IMAGE.
143 //EXPECT_EQ(0, SignChallenge(&message, signature, &siglen));
144 Reset(client, ResetRequest::PRODUCTION, signature, siglen);
145 bool bootloader;
146 bool production;
147 uint8_t locks[4];
148 GetState(client, &bootloader, &production, locks);
149 ASSERT_EQ(production, false);
150 }
151
152 } // namespace nugget_tools
153