From c7e9bb8c58c6fe1bc96ddf698fcaa17c03da3ca5 Mon Sep 17 00:00:00 2001
From: Olaoluwa Osuntokun <laolu32@gmail.com>
Date: Tue, 21 Jun 2016 11:46:27 -0700
Subject: [PATCH] lnrpc: implement wallet balance RPC call

---
 lnrpc/rpc.proto |  6 ++++++
 rpcserver.go    | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/lnrpc/rpc.proto b/lnrpc/rpc.proto
index 2ae96e08..7cbebff4 100644
--- a/lnrpc/rpc.proto
+++ b/lnrpc/rpc.proto
@@ -7,6 +7,7 @@ service Lightning {
     rpc NewAddress(NewAddressRequest) returns (NewAddressResponse);
 
     rpc ConnectPeer(ConnectPeerRequest) returns (ConnectPeerResponse);
+    rpc WalletBalance(WalletBalanceRequest) returns (WalletBalanceResponse);
 }
 
 message SendManyRequest {
@@ -36,4 +37,9 @@ message ConnectPeerRequest {
 
 message ConnectPeerResponse {
 	bytes lnID = 1;
+message WalletBalanceRequest {
+    bool witness_only = 1;
+}
+message WalletBalanceResponse {
+    double balance = 1;
 }
diff --git a/rpcserver.go b/rpcserver.go
index b6bd6647..614d8a93 100644
--- a/rpcserver.go
+++ b/rpcserver.go
@@ -134,4 +134,40 @@ func (r *rpcServer) ConnectPeer(ctx context.Context,
 
 	rpcsLog.Infof("Connected to peer: %v", peerAddr.String())
 	return &lnrpc.ConnectPeerResponse{[]byte(peerAddr.String())}, nil
+// WalletBalance returns the sum of all confirmed unspent outputs under control
+// by the wallet. This method can be modified by having the request specify
+// only witness outputs should be factored into the final output sum.
+// TODO(roasbeef): split into total and confirmed/unconfirmed
+func (r *rpcServer) WalletBalance(ctx context.Context,
+	in *lnrpc.WalletBalanceRequest) (*lnrpc.WalletBalanceResponse, error) {
+
+	var balance float64
+
+	if in.WitnessOnly {
+		witnessOutputs, err := r.server.lnwallet.ListUnspentWitness(1)
+		if err != nil {
+			return nil, err
+		}
+
+		// We need to convert from BTC to satoshi here otherwise, and
+		// incorrect sum will be returned.
+		var outputSum btcutil.Amount
+		for _, witnessOutput := range witnessOutputs {
+			outputSum += btcutil.Amount(witnessOutput.Amount * 1e8)
+		}
+
+		balance = outputSum.ToBTC()
+	} else {
+		// TODO(roasbeef): make num confs a param
+		outputSum, err := r.server.lnwallet.CalculateBalance(1)
+		if err != nil {
+			return nil, err
+		}
+
+		balance = outputSum.ToBTC()
+	}
+
+	rpcsLog.Debugf("walletbalance query response: %v", balance)
+
+	return &lnrpc.WalletBalanceResponse{balance}, nil
 }