github.com/hugh712/snapd@v0.0.0-20200910133618-1a99902bd583/overlord/state/copy.go (about)

     1  // -*- Mode: Go; indent-tabs-mode: t -*-
     2  
     3  /*
     4   * Copyright (C) 2020 Canonical Ltd
     5   *
     6   * This program is free software: you can redistribute it and/or modify
     7   * it under the terms of the GNU General Public License version 3 as
     8   * published by the Free Software Foundation.
     9   *
    10   * This program is distributed in the hope that it will be useful,
    11   * but WITHOUT ANY WARRANTY; without even the implied warranty of
    12   * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13   * GNU General Public License for more details.
    14   *
    15   * You should have received a copy of the GNU General Public License
    16   * along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17   *
    18   */
    19  
    20  package state
    21  
    22  import (
    23  	"bytes"
    24  	"encoding/json"
    25  	"fmt"
    26  	"os"
    27  	"path/filepath"
    28  	"strings"
    29  	"time"
    30  
    31  	"github.com/snapcore/snapd/jsonutil"
    32  	"github.com/snapcore/snapd/osutil"
    33  )
    34  
    35  type checkpointOnlyBackend struct {
    36  	path string
    37  }
    38  
    39  func (b *checkpointOnlyBackend) Checkpoint(data []byte) error {
    40  	if err := os.MkdirAll(filepath.Dir(b.path), 0755); err != nil {
    41  		return err
    42  	}
    43  	return osutil.AtomicWriteFile(b.path, data, 0600, 0)
    44  }
    45  
    46  func (b *checkpointOnlyBackend) EnsureBefore(d time.Duration) {
    47  	panic("cannot use EnsureBefore in checkpointOnlyBackend")
    48  }
    49  
    50  func (b *checkpointOnlyBackend) RequestRestart(t RestartType) {
    51  	panic("cannot use RequestRestart in checkpointOnlyBackend")
    52  }
    53  
    54  // copyData will copy the given subkeys specifier from srcData to dstData.
    55  //
    56  // The subkeys is constructed from a dotted path like "user.auth". This copy
    57  // helper is recursive and the pos parameter tells the function the current
    58  // position of the copy.
    59  func copyData(subkeys []string, pos int, srcData map[string]*json.RawMessage, dstData map[string]interface{}) error {
    60  	if pos < 0 || pos > len(subkeys) {
    61  		return fmt.Errorf("internal error: copyData used with an out-of-bounds position: %v not in [0:%v]", pos, len(subkeys))
    62  	}
    63  	raw, ok := srcData[subkeys[pos]]
    64  	if !ok {
    65  		return ErrNoState
    66  	}
    67  
    68  	if pos+1 == len(subkeys) {
    69  		dstData[subkeys[pos]] = raw
    70  		return nil
    71  	}
    72  
    73  	var srcDatam map[string]*json.RawMessage
    74  	if err := jsonutil.DecodeWithNumber(bytes.NewReader(*raw), &srcDatam); err != nil {
    75  		return fmt.Errorf("cannot unmarshal state entry %q with value %q as a map while trying to copy over %q", strings.Join(subkeys[:pos+1], "."), *raw, strings.Join(subkeys, "."))
    76  	}
    77  
    78  	// no subkey entry -> create one
    79  	if _, ok := dstData[subkeys[pos]]; !ok {
    80  		dstData[subkeys[pos]] = make(map[string]interface{})
    81  	}
    82  	// and use existing data
    83  	var dstDatam map[string]interface{}
    84  	switch dstDataEntry := dstData[subkeys[pos]].(type) {
    85  	case map[string]interface{}:
    86  		dstDatam = dstDataEntry
    87  	case *json.RawMessage:
    88  		dstDatam = make(map[string]interface{})
    89  		if err := jsonutil.DecodeWithNumber(bytes.NewReader(*dstDataEntry), &dstDatam); err != nil {
    90  			return fmt.Errorf("internal error: cannot decode subkey %s (%q) for %v (%T)", subkeys[pos], strings.Join(subkeys, "."), dstData, dstDataEntry)
    91  		}
    92  	default:
    93  		return fmt.Errorf("internal error: cannot create subkey %s (%q) for %v (%T)", subkeys[pos], strings.Join(subkeys, "."), dstData, dstData[subkeys[pos]])
    94  	}
    95  
    96  	return copyData(subkeys, pos+1, srcDatam, dstDatam)
    97  }
    98  
    99  // CopyState takes a state from the srcStatePath and copies all
   100  // dataEntries to the dstPath. Note that srcStatePath should never
   101  // point to a state that is in use.
   102  func CopyState(srcStatePath, dstStatePath string, dataEntries []string) error {
   103  	if osutil.FileExists(dstStatePath) {
   104  		// XXX: TOCTOU - look into moving this check into
   105  		// checkpointOnlyBackend. The issue is right now State
   106  		// will simply panic if Commit() returns an error
   107  		return fmt.Errorf("cannot copy state: %q already exists", dstStatePath)
   108  	}
   109  	if len(dataEntries) == 0 {
   110  		return fmt.Errorf("cannot copy state: must provide at least one data entry to copy")
   111  	}
   112  
   113  	f, err := os.Open(srcStatePath)
   114  	if err != nil {
   115  		return fmt.Errorf("cannot open state: %s", err)
   116  	}
   117  	defer f.Close()
   118  
   119  	// No need to lock/unlock the state here, srcState should not be
   120  	// in use at all.
   121  	srcState, err := ReadState(nil, f)
   122  	if err != nil {
   123  		return err
   124  	}
   125  
   126  	// copy relevant data
   127  	dstData := make(map[string]interface{})
   128  	for _, dataEntry := range dataEntries {
   129  		subkeys := strings.Split(dataEntry, ".")
   130  		if err := copyData(subkeys, 0, srcState.data, dstData); err != nil && err != ErrNoState {
   131  			return err
   132  		}
   133  	}
   134  
   135  	// write it out
   136  	dstState := New(&checkpointOnlyBackend{path: dstStatePath})
   137  	dstState.Lock()
   138  	defer dstState.Unlock()
   139  	for k, v := range dstData {
   140  		dstState.Set(k, v)
   141  	}
   142  
   143  	return nil
   144  }