Skip to content

Commit

Permalink
Add bindings support for static keys
Browse files Browse the repository at this point in the history 
Fixed: b/396642162

Change-Id: Ie0da0ee5052546f4980d489c554210fa912fc6b2
  • Loading branch information
@jblebrun
jblebrun committed Feb 18, 2025
1 parent 155f0b4 commit 23dae09
Show file tree
Hide file tree
Showing 5 changed files with 281 additions and 7 deletions.
8 changes: 8 additions & 0 deletions cc/oak_session/oak_session_bindings.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,7 @@ struct ErrorOrFfiEndorser {
};

struct SigningKey;
struct IdentityKey;

extern "C" {

Expand All @@ -167,6 +168,10 @@ extern SessionConfigBuilder* session_config_builder_add_peer_verifier(
SessionConfigBuilder*, BytesView, FfiAttestationVerifier);
extern SessionConfigBuilder* session_config_builder_add_session_binder(
SessionConfigBuilder*, BytesView, SigningKey*);
extern SessionConfigBuilder* session_config_builder_set_self_private_key(
SessionConfigBuilder*, IdentityKey*);
extern SessionConfigBuilder* session_config_builder_set_peer_static_public_key(
SessionConfigBuilder*, BytesView);
extern RustBytes new_fake_evidence(BytesView, BytesView);
extern ErrorOrFfiAttester new_simple_attester(BytesView);
extern RustBytes new_fake_endorsements(BytesView);
Expand All @@ -177,6 +182,9 @@ extern SigningKey* new_random_signing_key();
extern RustBytes signing_key_verifying_key_bytes(SigningKey*);
extern void free_signing_key(SigningKey*);

extern IdentityKey* new_identity_key();
extern ErrorOrRustBytes identity_key_get_public_key(IdentityKey*);

extern SessionConfig* session_config_builder_build(SessionConfigBuilder*);

// Corresponds to functions in oak_session/ffi/client_session.rs
Expand Down
101 changes: 100 additions & 1 deletion cc/oak_session/oak_session_bindings_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,56 @@ SessionConfig* TestConfigAttestedNNClient() {
return session_config_builder_build(session_config_builder);
}

TEST(OakSessionBindingsTest, TestHandshake) {
SessionConfig* TestConfigUnattestedNKClient(BytesView public_key) {
auto result = new_session_config_builder(ATTESTATION_TYPE_UNATTESTED,
HANDSHAKE_TYPE_NOISE_NK);
if (result.error != nullptr) {
LOG(FATAL) << "Failed to create session config builder"
<< static_cast<absl::string_view>(result.error->message);
}
auto session_config_builder = result.result;

session_config_builder = session_config_builder_set_peer_static_public_key(
session_config_builder, public_key);

return session_config_builder_build(session_config_builder);
}

SessionConfig* TestConfigUnattestedNKServer(IdentityKey* identity_key) {
auto result = new_session_config_builder(ATTESTATION_TYPE_UNATTESTED,
HANDSHAKE_TYPE_NOISE_NK);
if (result.error != nullptr) {
LOG(FATAL) << "Failed to create session config builder"
<< static_cast<absl::string_view>(result.error->message);
}
auto session_config_builder = result.result;

session_config_builder = session_config_builder_set_self_private_key(
session_config_builder, identity_key);

return session_config_builder_build(session_config_builder);
}

SessionConfig* TestConfigUnattestedKK(BytesView peer_public_key,
IdentityKey* self_identity_key) {
auto result = new_session_config_builder(ATTESTATION_TYPE_UNATTESTED,
HANDSHAKE_TYPE_NOISE_KK);
if (result.error != nullptr) {
LOG(FATAL) << "Failed to create session config builder"
<< static_cast<absl::string_view>(result.error->message);
}
auto session_config_builder = result.result;

session_config_builder = session_config_builder_set_self_private_key(
session_config_builder, self_identity_key);

session_config_builder = session_config_builder_set_peer_static_public_key(
session_config_builder, peer_public_key);

return session_config_builder_build(session_config_builder);
}

TEST(OakSessionBindingsTest, TestNNHandshake) {
ErrorOrServerSession server_session_result =
new_server_session(TestConfigUnattestedNN());
ASSERT_THAT(server_session_result, IsResult());
Expand All @@ -161,6 +210,56 @@ TEST(OakSessionBindingsTest, TestHandshake) {
free_client_session(client_session);
}

TEST(OakSessionBindingsTest, TestNKHandshake) {
IdentityKey* identity_key = new_identity_key();
ErrorOrRustBytes public_key = identity_key_get_public_key(identity_key);
ASSERT_THAT(public_key, IsResult());

ErrorOrServerSession server_session_result =
new_server_session(TestConfigUnattestedNKServer(identity_key));
ASSERT_THAT(server_session_result, IsResult());
ServerSession* server_session = server_session_result.result;
ErrorOrClientSession client_session_result = new_client_session(
TestConfigUnattestedNKClient(BytesView(*(public_key.result))));
ASSERT_THAT(client_session_result, IsResult());
ClientSession* client_session = client_session_result.result;

DoHandshake(server_session, client_session);

free_rust_bytes(public_key.result);
free_server_session(server_session);
free_client_session(client_session);
}

TEST(OakSessionBindingsTest, TestKKHandshake) {
IdentityKey* client_identity_key = new_identity_key();
IdentityKey* server_identity_key = new_identity_key();
ErrorOrRustBytes client_public_key =
identity_key_get_public_key(client_identity_key);
ASSERT_THAT(client_public_key, IsResult());
ErrorOrRustBytes server_public_key =
identity_key_get_public_key(server_identity_key);
ASSERT_THAT(server_public_key, IsResult());

ErrorOrServerSession server_session_result =
new_server_session(TestConfigUnattestedKK(
BytesView(*(client_public_key.result)), server_identity_key));
ASSERT_THAT(server_session_result, IsResult());
ServerSession* server_session = server_session_result.result;
ErrorOrClientSession client_session_result =
new_client_session(TestConfigUnattestedKK(
BytesView(*(server_public_key.result)), client_identity_key));
ASSERT_THAT(client_session_result, IsResult());
ClientSession* client_session = client_session_result.result;

DoHandshake(server_session, client_session);

free_rust_bytes(client_public_key.result);
free_rust_bytes(server_public_key.result);
free_server_session(server_session);
free_client_session(client_session);
}

TEST(OakSessionBindingsTest, AcceptEmptyOutgoingMessage) {
ErrorOrServerSession server_session_result =
new_server_session(TestConfigUnattestedNN());
Expand Down
1 change: 1 addition & 0 deletions oak_session/ffi/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,7 @@ rust_test(
":oak_session_ffi_server_session",
":oak_session_ffi_testing",
":oak_session_ffi_types",
"//oak_crypto",
"//oak_proto_rust",
"//oak_session",
"@oak_crates_index//:prost",
Expand Down
58 changes: 57 additions & 1 deletion oak_session/ffi/config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,13 +36,14 @@ use std::ffi::c_void;

use oak_attestation_types::{attester::Attester, endorser::Endorser};
use oak_attestation_verification_types::verifier::AttestationVerifier;
use oak_crypto::identity_key::{IdentityKey, IdentityKeyHandle};
use oak_session::{
attestation::AttestationType,
config::{SessionConfig, SessionConfigBuilder},
handshake::HandshakeType,
session_binding::SignatureBinderBuilder,
};
use oak_session_ffi_types::{BytesView, Error};
use oak_session_ffi_types::{BytesView, Error, ErrorOrRustBytes};
use p256::ecdsa::SigningKey;

/// Create a new `SessionConfigBuilder` for use in FFI;
Expand Down Expand Up @@ -252,6 +253,61 @@ pub unsafe extern "C" fn session_config_builder_add_session_binder(
Box::into_raw(Box::new(next_builder))
}

/// Call set_peer_static_public_key on the provided builder.
///
/// # Safety
///
/// * builder is a valid, properly aligned pointer to a SessionConfigBuilder.
/// * public_key is a valid, properly aligned BytesView.
#[no_mangle]
pub unsafe extern "C" fn session_config_builder_set_peer_static_public_key(
builder: *mut SessionConfigBuilder,
public_key: BytesView,
) -> *mut SessionConfigBuilder {
let next_builder = Box::from_raw(builder).set_peer_static_public_key(public_key.as_slice());
Box::into_raw(Box::new(next_builder))
}

/// Call set_self_private_key on the provided builder.
///
/// This function consumes the provided IdentityKey, reclaiming ownership on the
/// Rust side.
///
/// # Safety
///
/// * builder is a valid, properly aligned pointer to a SessionConfigBuilder.
/// * identity is a valid, properly aligned, acquired from a suitable source.
#[no_mangle]
pub unsafe extern "C" fn session_config_builder_set_self_private_key(
builder: *mut SessionConfigBuilder,
identity_key_ptr: *mut IdentityKey,
) -> *mut SessionConfigBuilder {
let identity_key = Box::from_raw(identity_key_ptr);
let next_builder = Box::from_raw(builder).set_self_private_key(identity_key);
Box::into_raw(Box::new(next_builder))
}

/// Create a new IdentityKey instance.
#[no_mangle]
pub extern "C" fn new_identity_key() -> *mut IdentityKey {
Box::into_raw(Box::new(IdentityKey::generate()))
}

/// Call get_public_key on the provided IdentityKey.
///
/// # Safety
///
/// * identity_key is valid, properly aligned, acquired from a suitable source.
#[no_mangle]
pub unsafe extern "C" fn identity_key_get_public_key(
identity_key: *const IdentityKey,
) -> ErrorOrRustBytes {
match (*identity_key).get_public_key() {
Ok(ik) => ErrorOrRustBytes::ok(ik.into_boxed_slice()),
Err(e) => ErrorOrRustBytes::err(e.to_string()),
}
}

#[repr(C)]
pub struct ErrorOrSessionConfigBuilder {
pub result: *mut SessionConfigBuilder,
Expand Down
120 changes: 115 additions & 5 deletions oak_session/ffi/tests.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
use oak_crypto::identity_key::IdentityKey;
use oak_proto_rust::oak::session::v1::{SessionRequest, SessionResponse};
use oak_session::{
session::{ClientSession, ServerSession},
Expand All @@ -20,16 +21,17 @@ use oak_session::{
use oak_session_ffi_attestation::{new_simple_attester, new_simple_endorser};
use oak_session_ffi_client_session as client_ffi;
use oak_session_ffi_config::{
self as config_ffi, session_config_builder_add_peer_verifier,
session_config_builder_add_self_attester, session_config_builder_add_self_endorser,
session_config_builder_add_session_binder,
self as config_ffi, identity_key_get_public_key, new_identity_key,
session_config_builder_add_peer_verifier, session_config_builder_add_self_attester,
session_config_builder_add_self_endorser, session_config_builder_add_session_binder,
session_config_builder_set_peer_static_public_key, session_config_builder_set_self_private_key,
};
use oak_session_ffi_server_session as server_ffi;
use oak_session_ffi_testing::{
free_signing_key, new_fake_attestation_verifier, new_fake_endorsements, new_fake_evidence,
new_random_signing_key, signing_key_verifying_key_bytes,
};
use oak_session_ffi_types::{free_rust_bytes_contents, BytesView};
use oak_session_ffi_types::{free_rust_bytes, free_rust_bytes_contents, BytesView};
use prost::Message;

macro_rules! assert_no_error {
Expand All @@ -44,7 +46,7 @@ macro_rules! assert_no_error {
}

#[test]
fn test_handshake() {
fn test_nn_handshake() {
let client_session_ptr = create_client_session(create_unattested_nn_session_config());
let server_session_ptr = create_server_session(create_unattested_nn_session_config());

Expand All @@ -56,6 +58,51 @@ fn test_handshake() {
unsafe { server_ffi::free_server_session(server_session_ptr) };
}

#[test]
fn test_nk_handshake() {
let identity_key = new_identity_key();
let client_session_ptr =
create_client_session(create_unattested_nk_client_session_config(identity_key));
let server_session_ptr =
create_server_session(create_unattested_nk_server_session_config(identity_key));

unsafe { do_handshake(client_session_ptr, server_session_ptr) };

assert!(unsafe { client_ffi::client_is_open(client_session_ptr) });
assert!(unsafe { server_ffi::server_is_open(server_session_ptr) });
unsafe { client_ffi::free_client_session(client_session_ptr) };
unsafe { server_ffi::free_server_session(server_session_ptr) };
}

#[test]
fn test_kk_handshake() {
let client_identity_key = new_identity_key();
let client_pk_result = unsafe { identity_key_get_public_key(client_identity_key) };
assert_no_error!(client_pk_result.error);

let server_identity_key = new_identity_key();
let server_pk_result = unsafe { identity_key_get_public_key(server_identity_key) };
assert_no_error!(server_pk_result.error);

let client_session_ptr = create_client_session(create_unattested_kk_session_config(
(unsafe { *server_pk_result.result }).as_bytes_view(),
client_identity_key,
));
let server_session_ptr = create_server_session(create_unattested_kk_session_config(
(unsafe { *client_pk_result.result }).as_bytes_view(),
server_identity_key,
));

unsafe { do_handshake(client_session_ptr, server_session_ptr) };

assert!(unsafe { client_ffi::client_is_open(client_session_ptr) });
assert!(unsafe { server_ffi::server_is_open(server_session_ptr) });
unsafe { free_rust_bytes(client_pk_result.result) };
unsafe { free_rust_bytes(server_pk_result.result) };
unsafe { client_ffi::free_client_session(client_session_ptr) };
unsafe { server_ffi::free_server_session(server_session_ptr) };
}

#[test]
fn test_client_encrypt_server_decrypt() {
let client_session_ptr = create_client_session(create_unattested_nn_session_config());
Expand Down Expand Up @@ -245,6 +292,69 @@ fn create_unattested_nn_session_config() -> *mut oak_session::config::SessionCon
unsafe { config_ffi::session_config_builder_build(session_config_builder.result) }
}

fn create_unattested_nk_client_session_config(
identity_key: *const IdentityKey,
) -> *mut oak_session::config::SessionConfig {
let session_config_builder_result = config_ffi::new_session_config_builder(
config_ffi::ATTESTATION_TYPE_UNATTESTED,
config_ffi::HANDSHAKE_TYPE_NOISE_NK,
);
assert_no_error!(session_config_builder_result.error);
let mut session_config_builder = session_config_builder_result.result;

let public_key_result = unsafe { identity_key_get_public_key(identity_key) };
assert_no_error!(public_key_result.error);

session_config_builder = unsafe {
session_config_builder_set_peer_static_public_key(
session_config_builder,
(*public_key_result.result).as_bytes_view(),
)
};

unsafe { free_rust_bytes(public_key_result.result) }
unsafe { config_ffi::session_config_builder_build(session_config_builder) }
}

fn create_unattested_nk_server_session_config(
identity_key: *mut IdentityKey,
) -> *mut oak_session::config::SessionConfig {
let session_config_builder_result = config_ffi::new_session_config_builder(
config_ffi::ATTESTATION_TYPE_UNATTESTED,
config_ffi::HANDSHAKE_TYPE_NOISE_NK,
);
assert_no_error!(session_config_builder_result.error);
let mut session_config_builder = session_config_builder_result.result;

session_config_builder = unsafe {
session_config_builder_set_self_private_key(session_config_builder, identity_key)
};

unsafe { config_ffi::session_config_builder_build(session_config_builder) }
}

fn create_unattested_kk_session_config(
peer_public_key: BytesView,
self_identity_key: *mut IdentityKey,
) -> *mut oak_session::config::SessionConfig {
let session_config_builder_result = config_ffi::new_session_config_builder(
config_ffi::ATTESTATION_TYPE_UNATTESTED,
config_ffi::HANDSHAKE_TYPE_NOISE_KK,
);
assert_no_error!(session_config_builder_result.error);
let mut session_config_builder = session_config_builder_result.result;

session_config_builder = unsafe {
session_config_builder_set_self_private_key(session_config_builder, self_identity_key)
};

session_config_builder = unsafe {
session_config_builder_set_peer_static_public_key(session_config_builder, peer_public_key)
};

unsafe { config_ffi::session_config_builder_build(session_config_builder) }
}

fn create_attested_nn_server_session_config() -> *mut oak_session::config::SessionConfig {
let session_config_builder_result = config_ffi::new_session_config_builder(
config_ffi::ATTESTATION_TYPE_SELF_UNIDIRECTIONAL,
Expand Down

0 comments on commit 23dae09

Please sign in to comment.