| // Copyright 2022 The ChromiumOS Authors |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| package dut |
| |
| import ( |
| "compress/gzip" |
| "context" |
| "fmt" |
| "io" |
| "log" |
| "os" |
| "path" |
| "path/filepath" |
| "strconv" |
| "syscall" |
| |
| "cloud.google.com/go/storage" |
| "github.com/klauspost/compress/zstd" |
| "github.com/klauspost/readahead" |
| "golang.org/x/oauth2" |
| "google.golang.org/api/option" |
| |
| "chromium.googlesource.com/chromiumos/platform/dev-util.git/contrib/fflash/internal/artifact" |
| "chromium.googlesource.com/chromiumos/platform/dev-util.git/contrib/fflash/internal/misc" |
| "chromium.googlesource.com/chromiumos/platform/dev-util.git/contrib/fflash/internal/progress" |
| ) |
| |
| // copyChunked copies r to w in chunks. |
| func copyChunked(w io.Writer, r io.Reader, buf []byte) (written int64, err error) { |
| for { |
| n, err := io.ReadFull(r, buf) |
| if err != nil && err != io.ErrUnexpectedEOF { |
| if err == io.EOF { |
| break |
| } |
| return written, err |
| } |
| if m, err := w.Write(buf[:n]); err != nil { |
| return written, err |
| } else { |
| written += int64(m) |
| } |
| } |
| return written, nil |
| } |
| |
| type closeFunc func() error |
| |
| // Client creates a storage.Client from req. |
| func (req *Request) Client(ctx context.Context) (*storage.Client, error) { |
| client, err := storage.NewClient(ctx, |
| option.WithTokenSource(oauth2.StaticTokenSource(req.Token)), |
| ) |
| if err != nil { |
| return nil, fmt.Errorf("storage.NewClient failed: %w", err) |
| } |
| |
| return client, nil |
| } |
| |
| // objectHandle returns the storage.ObjectHandle for the artifact bucket, dir and name. |
| func objectHandle(client *storage.Client, bucket, dir, name string) *storage.ObjectHandle { |
| return client.Bucket(bucket).Object(path.Join(dir, name)) |
| } |
| |
| // openObject opens the file in the the directory specified by req. |
| func (req *Request) openObject(ctx context.Context, client *storage.Client, rw *progress.ReportingWriter, name string, decompress bool) (io.Reader, closeFunc, error) { |
| obj := objectHandle(client, req.Artifact.Bucket, req.Artifact.Dir, name) |
| |
| rd, err := obj.NewReader(ctx) |
| if err != nil { |
| return nil, nil, fmt.Errorf("obj.NewReader for %s failed: %w", misc.GsURI(obj), err) |
| } |
| rw.SetTotal(rd.Attrs.Size) |
| |
| aheadRd, err := readahead.NewReadCloserSize(rd, 4, 1<<20) |
| if err != nil { |
| rd.Close() |
| return nil, nil, fmt.Errorf( |
| "readahead.NewReadCloserSize for %s failed: %w", misc.GsURI(obj), err) |
| } |
| |
| brd := io.TeeReader(aheadRd, rw) |
| |
| var decompressRd io.ReadCloser |
| if decompress { |
| switch req.Artifact.Compression { |
| case artifact.Gzip: |
| decompressRd, err = gzip.NewReader(brd) |
| if err != nil { |
| rd.Close() |
| return nil, nil, fmt.Errorf("gzip.NewReader for %s failed: %w", misc.GsURI(obj), err) |
| } |
| case artifact.Zstd: |
| var d *zstd.Decoder |
| d, err = zstd.NewReader(brd) |
| if err != nil { |
| rd.Close() |
| return nil, nil, fmt.Errorf("zstd.NewReader for %s failed: %w", misc.GsURI(obj), err) |
| } |
| decompressRd = d.IOReadCloser() |
| default: |
| log.Panicf("unknown compression %q", req.Artifact.Compression) |
| } |
| } else { |
| decompressRd = io.NopCloser(brd) |
| } |
| |
| return decompressRd, |
| func() error { |
| decompressRd.Close() |
| return aheadRd.Close() |
| }, |
| nil |
| } |
| |
| // CheckArtifact checks that all image artifacts are accessible. |
| func CheckArtifact(ctx context.Context, client *storage.Client, artifact *artifact.Artifact) error { |
| for _, file := range artifact.Images.Names() { |
| obj := objectHandle(client, artifact.Bucket, artifact.Dir, file) |
| |
| if _, err := obj.Attrs(ctx); err != nil { |
| return fmt.Errorf("%s: %w", misc.GsURI(obj), err) |
| } |
| } |
| return nil |
| } |
| |
| // Flash a partition with imageGz to partition. |
| func (req *Request) Flash(ctx context.Context, client *storage.Client, rw *progress.ReportingWriter, image string, partition string) error { |
| r, close, err := req.openObject(ctx, client, rw, image, true) |
| if err != nil { |
| return err |
| } |
| defer close() |
| |
| w, err := os.OpenFile(partition, os.O_WRONLY|syscall.O_DIRECT, 0660) |
| if err != nil { |
| return fmt.Errorf("cannot open %s: %w", partition, err) |
| } |
| defer func() { |
| if err := w.Close(); err != nil { |
| panic(err) |
| } |
| }() |
| |
| if _, err := copyChunked(w, r, make([]byte, 1<<20)); err != nil { |
| return fmt.Errorf("copy to %s failed: %w", partition, err) |
| } |
| |
| return nil |
| } |
| |
| // FlashStateful flashes the stateful partition. |
| func (req *Request) FlashStateful(ctx context.Context, client *storage.Client, rw *progress.ReportingWriter, clobber bool) error { |
| r, close, err := req.openObject(ctx, client, rw, req.Artifact.Stateful, false) |
| if err != nil { |
| return err |
| } |
| defer close() |
| |
| if err := unpackStateful(ctx, r, req.Artifact.Compression); err != nil { |
| return err |
| } |
| |
| content := "" |
| if clobber { |
| content = "clobber" |
| } |
| |
| if err := os.WriteFile( |
| filepath.Join(statefulDir, statefulAvailable), |
| []byte(content), |
| 0644, |
| ); err != nil { |
| return fmt.Errorf("failed to write %s: %w", statefulAvailable, err) |
| } |
| |
| return nil |
| } |
| |
| // RunPostinst runs "postinst" from the partition. |
| func RunPostinst(ctx context.Context, partition string) error { |
| dir, err := os.MkdirTemp("/tmp", "dut-agent-*") |
| if err != nil { |
| return err |
| } |
| |
| if _, err := runCommand(ctx, "mount", "-o", "ro", partition, dir); err != nil { |
| return err |
| } |
| defer func() { |
| if _, err := runCommand(context.Background(), "umount", partition); err != nil { |
| log.Printf("failed to unmount rootfs: %s", err) |
| } |
| }() |
| |
| return runCommandStderr(ctx, filepath.Join(dir, "postinst"), partition) |
| } |
| |
| func DisableRootfsVerification(ctx context.Context, kernelNum int) error { |
| return runCommandStderr(ctx, |
| "/usr/share/vboot/bin/make_dev_ssd.sh", |
| "--remove_rootfs_verification", |
| "--partitions", strconv.Itoa(kernelNum), |
| ) |
| } |
| |
| func ClearTpmOwner(ctx context.Context) error { |
| return runCommandStderr(ctx, |
| "crossystem", |
| "clear_tpm_owner_request=1", |
| ) |
| } |