1// Copyright (c) 2012 The Chromium OS Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5#include <string.h>
6
7#define _STUB_IMPLEMENTATION_
8
9#include "cgpt.h"
10#include "cgpt_params.h"
11#include "cgptlib_internal.h"
12#include "utility.h"
13#include "vboot_host.h"
14
15static const char* DumpCgptAddParams(const CgptAddParams *params) {
16  static char buf[256];
17  char tmp[64];
18
19  buf[0] = 0;
20  snprintf(tmp, sizeof(tmp), "-i %d ", params->partition);
21  StrnAppend(buf, tmp, sizeof(buf));
22  if (params->label) {
23    snprintf(tmp, sizeof(tmp), "-l %s ", params->label);
24    StrnAppend(buf, tmp, sizeof(buf));
25  }
26  if (params->set_begin) {
27    snprintf(tmp, sizeof(tmp), "-b %llu ", (unsigned long long)params->begin);
28    StrnAppend(buf, tmp, sizeof(buf));
29  }
30  if (params->set_size) {
31    snprintf(tmp, sizeof(tmp), "-s %llu ", (unsigned long long)params->size);
32    StrnAppend(buf, tmp, sizeof(buf));
33  }
34  if (params->set_type) {
35    GuidToStr(&params->type_guid, tmp, sizeof(tmp));
36    StrnAppend(buf, "-t ", sizeof(buf));
37    StrnAppend(buf, tmp, sizeof(buf));
38    StrnAppend(buf, " ", sizeof(buf));
39  }
40  if (params->set_unique) {
41    GuidToStr(&params->unique_guid, tmp, sizeof(tmp));
42    StrnAppend(buf, "-u ", sizeof(buf));
43    StrnAppend(buf, tmp, sizeof(buf));
44    StrnAppend(buf, " ", sizeof(buf));
45  }
46  if (params->set_successful) {
47    snprintf(tmp, sizeof(tmp), "-S %d ", params->successful);
48    StrnAppend(buf, tmp, sizeof(buf));
49  }
50  if (params->set_tries) {
51    snprintf(tmp, sizeof(tmp), "-T %d ", params->tries);
52    StrnAppend(buf, tmp, sizeof(buf));
53  }
54  if (params->set_priority) {
55    snprintf(tmp, sizeof(tmp), "-P %d ", params->priority);
56    StrnAppend(buf, tmp, sizeof(buf));
57  }
58  if (params->set_raw) {
59    snprintf(tmp, sizeof(tmp), "-A 0x%x ", params->raw_value);
60    StrnAppend(buf, tmp, sizeof(buf));
61  }
62
63  StrnAppend(buf, "\n", sizeof(buf));
64  return buf;
65}
66
67// This is the implementation-specific helper function.
68static int GptSetEntryAttributes(struct drive *drive,
69                                 uint32_t index,
70                                 CgptAddParams *params) {
71  GptEntry *entry;
72
73  entry = GetEntry(&drive->gpt, PRIMARY, index);
74  if (params->set_begin)
75    entry->starting_lba = params->begin;
76  if (params->set_size)
77    entry->ending_lba = entry->starting_lba + params->size - 1;
78  if (params->set_unique) {
79    memcpy(&entry->unique, &params->unique_guid, sizeof(Guid));
80  } else if (GuidIsZero(&entry->type)) {
81	  if (CGPT_OK != GenerateGuid(&entry->unique)) {
82		  Error("Unable to generate new GUID.\n");
83		  return -1;
84    }
85  }
86  if (params->set_type)
87    memcpy(&entry->type, &params->type_guid, sizeof(Guid));
88  if (params->label) {
89    if (CGPT_OK != UTF8ToUTF16((uint8_t *)params->label, entry->name,
90                               sizeof(entry->name) / sizeof(entry->name[0]))) {
91      Error("The label cannot be converted to UTF16.\n");
92      return -1;
93    }
94  }
95  return 0;
96}
97
98// This is an internal helper function which assumes no NULL args are passed.
99// It sets the given attribute values for a single entry at the given index.
100static int SetEntryAttributes(struct drive *drive,
101                              uint32_t index,
102                              CgptAddParams *params) {
103  if (params->set_raw) {
104    SetRaw(drive, PRIMARY, index, params->raw_value);
105  } else {
106    if (params->set_successful)
107      SetSuccessful(drive, PRIMARY, index, params->successful);
108    if (params->set_tries)
109      SetTries(drive, PRIMARY, index, params->tries);
110    if (params->set_priority)
111      SetPriority(drive, PRIMARY, index, params->priority);
112  }
113
114  // New partitions must specify type, begin, and size.
115  if (IsUnused(drive, PRIMARY, index)) {
116    if (!params->set_begin || !params->set_size || !params->set_type) {
117      Error("-t, -b, and -s options are required for new partitions\n");
118      return -1;
119    }
120    if (GuidIsZero(&params->type_guid)) {
121      Error("New partitions must have a type other than \"unused\"\n");
122      return -1;
123    }
124  }
125
126  return 0;
127}
128
129static int CgptCheckAddValidity(struct drive *drive) {
130  int gpt_retval;
131  if (GPT_SUCCESS != (gpt_retval = GptSanityCheck(&drive->gpt))) {
132    Error("GptSanityCheck() returned %d: %s\n",
133          gpt_retval, GptError(gpt_retval));
134    return -1;
135  }
136
137  if (((drive->gpt.valid_headers & MASK_BOTH) != MASK_BOTH) ||
138      ((drive->gpt.valid_entries & MASK_BOTH) != MASK_BOTH)) {
139    Error("one of the GPT header/entries is invalid.\n"
140          "please run 'cgpt repair' before adding anything.\n");
141    return -1;
142  }
143
144  return 0;
145}
146
147static int CgptGetUnusedPartition(struct drive *drive, uint32_t *index,
148                                  CgptAddParams *params) {
149  uint32_t i;
150  uint32_t max_part = GetNumberOfEntries(drive);
151  if (params->partition) {
152    if (params->partition > max_part) {
153      Error("invalid partition number: %d\n", params->partition);
154      return -1;
155    }
156    *index = params->partition - 1;
157    return 0;
158  } else {
159    // Find next empty partition.
160    for (i = 0; i < max_part; i++) {
161      if (IsUnused(drive, PRIMARY, i)) {
162        params->partition = i + 1;
163        *index = i;
164        return 0;
165      }
166    }
167    Error("no unused partitions available\n");
168    return -1;
169  }
170}
171
172int CgptSetAttributes(CgptAddParams *params) {
173  struct drive drive;
174
175  if (params == NULL)
176    return CGPT_FAILED;
177
178  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
179                           params->drive_size))
180    return CGPT_FAILED;
181
182  if (CgptCheckAddValidity(&drive)) {
183    goto bad;
184  }
185
186  if (params->partition == 0 ||
187      params->partition >= GetNumberOfEntries(&drive)) {
188    Error("invalid partition number: %d\n", params->partition);
189    goto bad;
190  }
191
192  SetEntryAttributes(&drive, params->partition - 1, params);
193
194  UpdateAllEntries(&drive);
195
196  // Write it all out.
197  return DriveClose(&drive, 1);
198
199bad:
200  DriveClose(&drive, 0);
201  return CGPT_FAILED;
202}
203
204// This method gets the partition details such as the attributes, the
205// guids of the partitions, etc. Input is the partition number or the
206// unique id of the partition. Output is populated in the respective
207// fields of params.
208int CgptGetPartitionDetails(CgptAddParams *params) {
209  struct drive drive;
210  int result = CGPT_FAILED;
211  int index;
212
213  if (params == NULL)
214    return CGPT_FAILED;
215
216  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
217                           params->drive_size))
218    return CGPT_FAILED;
219
220  if (CgptCheckAddValidity(&drive)) {
221    goto bad;
222  }
223
224  int max_part = GetNumberOfEntries(&drive);
225  if (params->partition > 0) {
226    if (params->partition >= max_part) {
227      Error("invalid partition number: %d\n", params->partition);
228      goto bad;
229    }
230  } else {
231    if (!params->set_unique) {
232      Error("either partition or unique_id must be specified\n");
233      goto bad;
234    }
235    for (index = 0; index < max_part; index++) {
236      GptEntry *entry = GetEntry(&drive.gpt, PRIMARY, index);
237      if (GuidEqual(&entry->unique, &params->unique_guid)) {
238        params->partition = index + 1;
239        break;
240      }
241    }
242    if (index >= max_part) {
243      Error("no partitions with the given unique id available\n");
244      goto bad;
245    }
246  }
247  index = params->partition - 1;
248
249  // GPT-specific code
250  GptEntry *entry = GetEntry(&drive.gpt, PRIMARY, index);
251  params->begin = entry->starting_lba;
252  params->size =  entry->ending_lba - entry->starting_lba + 1;
253  memcpy(&params->type_guid, &entry->type, sizeof(Guid));
254  memcpy(&params->unique_guid, &entry->unique, sizeof(Guid));
255  params->raw_value = entry->attrs.fields.gpt_att;
256
257  params->successful = GetSuccessful(&drive, PRIMARY, index);
258  params->tries = GetTries(&drive, PRIMARY, index);
259  params->priority = GetPriority(&drive, PRIMARY, index);
260  result = CGPT_OK;
261
262bad:
263  DriveClose(&drive, 0);
264  return result;
265}
266
267static int GptAdd(struct drive *drive, CgptAddParams *params, uint32_t index) {
268  GptEntry *entry, backup;
269  int rv;
270
271  entry = GetEntry(&drive->gpt, PRIMARY, index);
272  memcpy(&backup, entry, sizeof(backup));
273
274  if (SetEntryAttributes(drive, index, params) ||
275      GptSetEntryAttributes(drive, index, params)) {
276    memcpy(entry, &backup, sizeof(*entry));
277    return -1;
278  }
279
280  UpdateAllEntries(drive);
281
282  rv = CheckEntries((GptEntry*)drive->gpt.primary_entries,
283                    (GptHeader*)drive->gpt.primary_header);
284
285  if (0 != rv) {
286    // If the modified entry is illegal, recover it and return error.
287    memcpy(entry, &backup, sizeof(*entry));
288    Error("%s\n", GptErrorText(rv));
289    Error(DumpCgptAddParams(params));
290    return -1;
291  }
292
293  return 0;
294}
295
296int CgptAdd(CgptAddParams *params) {
297  struct drive drive;
298  uint32_t index;
299
300  if (params == NULL)
301    return CGPT_FAILED;
302
303  if (CGPT_OK != DriveOpen(params->drive_name, &drive, O_RDWR,
304                           params->drive_size))
305    return CGPT_FAILED;
306
307  if (CgptCheckAddValidity(&drive)) {
308    goto bad;
309  }
310
311  if (CgptGetUnusedPartition(&drive, &index, params)) {
312    goto bad;
313  }
314
315  if (GptAdd(&drive, params, index))
316    goto bad;
317
318  // Write it all out.
319  return DriveClose(&drive, 1);
320
321bad:
322  DriveClose(&drive, 0);
323  return CGPT_FAILED;
324}
325