Golang MessageStream.ReadMessage方法代码示例

// RunMessageLoop handles messages from a client on a given message stream until
// it gets an error trying to read a message.
func (m *ResourceMaster) RunMessageLoop(ms util.MessageStream, programPolicy *ProgramPolicy) error {
	for {
		var msg Message
		if err := ms.ReadMessage(&msg); err != nil {
			return err
		}

		var fop FileOperation
		t := *msg.Type
		if t == MessageType_CREATE || t == MessageType_DELETE || t == MessageType_READ || t == MessageType_WRITE {
			if err := proto.Unmarshal(msg.Data, &fop); err != nil {
				log.Printf("Couldn't unmarshal FileOperation for operation %d\n", t)
				continue
			}

			if err := m.checkFileAuth(&msg, &fop); err != nil {
				log.Printf("The file operation %d didn't pass authorization: %s\n", t, err)
				continue
			}
		}

		switch *msg.Type {
		case MessageType_AUTH_CERT:
			cert, err := m.AuthenticatePrincipal(ms, &msg, programPolicy)
			if err != nil {
				log.Printf("Failed to authenticate a principal: %s\n", err)
				continue
			}

			owner, err := principalNameFromDERCert(cert)
			if err != nil {
				log.Printf("Couldn't get the owner name from the cert: %s\n", err)
				continue
			}
			_, err = m.InsertPrincipal(owner, cert, Authenticated)
			if err != nil {
				log.Printf("Couldn't set the principal as authenticated")
			}
		case MessageType_CREATE:
			if err := m.Create(ms, &fop); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		case MessageType_READ:
			if err := m.Read(ms, &fop, programPolicy.SymKeys); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		case MessageType_WRITE:
			if err := m.Write(ms, &fop, programPolicy.SymKeys); err != nil {
				log.Printf("Couldn't create the file %s: %s\n", *fop.Name, err)
			}
		default:
			if err := sendResult(ms, false); err != nil {
				log.Printf("Couldn't signal failure for the invalid operation: %s", err)
			}
			log.Printf("Invalid initial message type %d\n", *msg.Type)
		}
	}
}

// checkResponse waits for a RollbackResponse and checks to make sure it's not
// an ERROR response from the server.
func checkResponse(ms *util.MessageStream) error {
	var rr RollbackResponse
	if err := ms.ReadMessage(&rr); err != nil {
		return err
	}
	if *rr.Type == RollbackMessageType_ERROR {
		return fmt.Errorf("couldn't set the counter on the remote server")
	}
	return nil
}

// recvResult waits for a OperationResult on a MessageStream
func recvResult(ms util.MessageStream) (bool, error) {
	var m Message
	if err := ms.ReadMessage(&m); err != nil {
		return false, err
	}
	var res OperationResult
	if err := proto.Unmarshal(m.Data, &res); err != nil {
		return false, err
	}

	return *res.Result, nil
}

func (conn *Conn) recvCredentials(ms util.MessageStream) {
	m, err := ms.ReadString()
	if err != nil {
		return
	}
	if m == "delegation" {
		var a Attestation
		if err = ms.ReadMessage(&a); err != nil {
			return
		}
		// Validate the peer certificate
		peerCert := conn.ConnectionState().PeerCertificates[0]
		p, err := ValidatePeerAttestation(&a, peerCert)
		if err != nil {
			ms.SetErr(err)
			return
		}
		if conn.guard != nil {
			if conn.verifier != nil {
				if err = AddEndorsements(conn.guard, &a, conn.verifier); err != nil {
					ms.SetErr(err)
					return
				}
			}
			if !conn.guard.IsAuthorized(p, "Connect", nil) {
				ms.SetErr(errors.New("principal delegator in client attestation is not authorized to connect"))
				return
			}
		}
		conn.peer = &p
	} else if m == "key" {
		peerCert := conn.ConnectionState().PeerCertificates[0]
		v, err := FromX509(peerCert)
		if err != nil {
			ms.SetErr(errors.New("can't decode key from peer certificate"))
			return
		}
		p := v.ToPrincipal()
		conn.peer = &p
	} else if m == "anonymous" {
		if conn.guard != nil {
			err = errors.New("peer did not provide tao delegation")
			ms.SetErr(err)
			return
		}
	} else {
		err = errors.New("unrecognized authentication handshake: " + m)
		ms.SetErr(err)
		return
	}
}

// readResult reads an OperationResult and returns its value or an error.
func readResult(ms util.MessageStream) (bool, error) {
	// Read the response wrapper message.
	var arm Message
	if err := ms.ReadMessage(&arm); err != nil {
		return false, err
	}
	if *arm.Type != MessageType_OP_RES {
		return false, fmt.Errorf("didn't receive OP_RES from the server")
	}

	var opr OperationResult
	if err := proto.Unmarshal(arm.Data, &opr); err != nil {
		return false, err
	}
	return *opr.Result, nil
}

// RunMessageLoop handles incoming messages for the RollbackMaster and passes
// them to the appropriate functions.
func (r *RollbackMaster) RunMessageLoop(ms *util.MessageStream, programPolicy *ProgramPolicy, name string) error {
	for {
		var msg RollbackMessage
		if err := ms.ReadMessage(&msg); err != nil {
			return err
		}

		switch *msg.Type {
		case RollbackMessageType_SET_COUNTER:
			i, err := decodeCounter(msg.Data)
			if err != nil {
				log.Printf("failed to decode counter for SET_COUNTER: %s", err)
				continue
			}

			if err = r.SetCounter(ms, name, i); err != nil {
				log.Printf("failed to set the counter on the RollbackMaster: %s", err)
				continue
			}
		case RollbackMessageType_GET_COUNTER:
			if err := r.GetCounter(ms, name); err != nil {
				log.Printf("failed to get the counter for program %s", name)
				continue
			}
		case RollbackMessageType_SET_HASH:
			var rh RollbackHash
			if err := proto.Unmarshal(msg.Data, &rh); err != nil {
				log.Printf("failed to unmarshal the parameters for SET_HASH: %s", err)
				continue
			}
			if err := r.SetHash(ms, name, *rh.Item, rh.Hash); err != nil {
				log.Printf("failed to set the hash for item %s on program %s: %s", *rh.Item, name, err)
				continue
			}
		case RollbackMessageType_GET_HASHED_VERIFIER:
			if err := r.GetHashedVerifier(ms, name, string(msg.Data)); err != nil {
				log.Printf("failed to get the hashed verifier for program %s: %s", name, err)
				continue
			}
		default:
			log.Printf("unknown rollback message %d", *msg.Type)
		}
	}
}

// AuthenticatePrincipal is a client method used to send a request to a
// ResourceMaster to authenticate a principal with a given certificate and a
// given set of keys.
func AuthenticatePrincipal(ms util.MessageStream, key *tao.Keys, derCert []byte) error {
	// Send the authentication request, which supposes that a server is
	// waiting to receive this request.
	c := &Message{
		Type: MessageType_AUTH_CERT.Enum(),
		Data: derCert,
	}
	if _, err := ms.WriteMessage(c); err != nil {
		return err
	}

	// Receive a challenge nonce from the server.
	var nc Message
	if err := ms.ReadMessage(&nc); err != nil {
		return err
	}
	if *nc.Type != MessageType_NONCE_CHALL {
		return fmt.Errorf("didn't receive NONCE_CHALL from the server")
	}

	// Sign the nonce.
	sn := &Message{
		Type: MessageType_SIGNED_NONCE.Enum(),
	}
	var err error
	if sn.Data, err = key.SigningKey.Sign(nc.Data, ChallengeContext); err != nil {
		return err
	}
	if _, err := ms.WriteMessage(sn); err != nil {
		return err
	}

	// Get the result from the server after verificaton.
	res, err := readResult(ms)
	if err != nil {
		return err
	}

	if !res {
		return fmt.Errorf("the signed nonce failed verification")
	}
	return nil
}

// GetHashedVerifier gets the hash of the counter and the item hash for a given
// item.
func GetHashedVerifier(ms *util.MessageStream, item string) ([]byte, error) {
	rm := &RollbackMessage{
		Type: RollbackMessageType_GET_HASHED_VERIFIER.Enum(),
		Data: []byte(item),
	}
	if _, err := ms.WriteMessage(rm); err != nil {
		return nil, err
	}

	// We can't use checkResponse here since we need to get the value out of
	// the response to read the hash.
	var rr RollbackResponse
	if err := ms.ReadMessage(&rr); err != nil {
		return nil, err
	}
	if *rr.Type == RollbackMessageType_ERROR {
		return nil, fmt.Errorf("couldn't set the counter on the remote server")
	}

	return rr.Data, nil
}

// GetCounter gets the current value of the monotonic counter for a given
// program name.
func GetCounter(ms *util.MessageStream) (uint64, error) {
	// The name of the program is managed by the rollback server, not the
	// client, so it doesn't need to be passed in the message.
	rm := &RollbackMessage{
		Type: RollbackMessageType_GET_COUNTER.Enum(),
		Data: make([]byte, 0),
	}
	if _, err := ms.WriteMessage(rm); err != nil {
		return 0, err
	}

	// We can't use checkResponse here since we need to get the value out of
	// the response to read the counter.
	var rr RollbackResponse
	if err := ms.ReadMessage(&rr); err != nil {
		return 0, err
	}
	if *rr.Type == RollbackMessageType_ERROR {
		return 0, fmt.Errorf("couldn't set the counter on the remote server")
	}

	return decodeCounter(rr.Data)
}

// GetFile receives bytes from a sender and optionally encrypts them and adds
// integrity protection, and writes them to disk.
func GetFile(ms util.MessageStream, dir string, filename string, keys []byte) error {
	fullpath := path.Join(dir, filename)
	file, err := os.Create(fullpath)
	if err != nil {
		return fmt.Errorf("can't create file '%s' in GetFile", fullpath)
	}
	defer file.Close()

	var ctr cipher.Stream
	var hm hash.Hash
	iv := make([]byte, ivSize)

	hasKeys := len(keys) >= minKeySize
	if hasKeys {
		enc, err := aes.NewCipher(keys[:aesKeySize])
		if err != nil || enc == nil {
			return fmt.Errorf("couldn't create an AES cipher: %s", err)
		}

		// Use the remaining bytes of the key slice for the HMAC key.
		hm = hmac.New(sha256.New, keys[aesKeySize:])
		if _, err := rand.Read(iv); err != nil {
			return fmt.Errorf("couldn't read random bytes for a fresh IV: %s", err)
		}

		// The first bytes of the HMAC input are the IV.
		hm.Write(iv)
		ctr = cipher.NewCTR(enc, iv)
		if ctr == nil {
			return fmt.Errorf("couldn't create a new instance of AES-CTR-128")
		}
		if _, err = file.Write(iv); err != nil {
			return err
		}
	}

	// temp holds temporary encrypted ciphertext before it's written to
	// disk.
	temp := make([]byte, bufferSize)
	for {
		var m Message
		if err := ms.ReadMessage(&m); err != nil {
			return nil
		}

		// Sanity check: this must be FILE_LAST or FILE_NEXT.
		t := *m.Type
		if !(t == MessageType_FILE_LAST || t == MessageType_FILE_NEXT) {
			return fmt.Errorf("received invalid message type %d during file streaming in GetFile", t)
		}

		if hasKeys {
			l := len(m.Data)
			ctr.XORKeyStream(temp, m.Data)
			hm.Write(temp[:l])
			if _, err = file.Write(temp[:l]); err != nil {
				return err
			}
		} else {
			if _, err = file.Write(m.Data); err != nil {
				return err
			}
		}

		// FILE_LAST corresponds to receiving the final bytes of the
		// file.
		if *m.Type == MessageType_FILE_LAST {
			break
		}
	}

	// Write the MAC at the end of the file.
	if hasKeys {
		hmacBytes := hm.Sum(nil)
		if _, err = file.Write(hmacBytes[:]); err != nil {
			return err
		}
	}

	return nil
}