diff --git a/v0/test_token.c b/v0/test_token.c
index 265d277..ca90182 100644
--- a/v0/test_token.c
+++ b/v0/test_token.c
@@ -33,14 +33,20 @@ static void test_token_to_string_identifier(void) {
 }
 
 static void test_tokenstream_open_fail(void) {
-	TokenStream* ts = tokenstream_open("v0/does_not_exist.c2");
-	if (ts != NULL) fail("expected NULL for non-existent file");
+	Buffer* buf = buffer_open_file("v0/does_not_exist.c2");
+	if (buf != NULL) fail("expected NULL for non-existent file");
+	
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts != NULL) fail("expected NULL for NULL buffer");
 }
 
 static void test_tokenstream_simple_keyword(void) {
 	write_test_file("v0/test_token_tmp.c2", "module");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	Token t = tokenstream_next(ts);
 	if (t != TOKEN_MODULE) fail("expected TOKEN_MODULE");
@@ -53,8 +59,11 @@ static void test_tokenstream_simple_keyword(void) {
 
 static void test_tokenstream_keywords_and_symbols(void) {
 	write_test_file("v0/test_token_tmp.c2", "module main; import stdio;");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	if (tokenstream_next(ts) != TOKEN_MODULE) fail("expected TOKEN_MODULE");
 	if (tokenstream_next(ts) != TOKEN_IDENTIFIER) fail("expected TOKEN_IDENTIFIER (main)");
@@ -69,8 +78,11 @@ static void test_tokenstream_keywords_and_symbols(void) {
 
 static void test_tokenstream_parentheses_and_brackets(void) {
 	write_test_file("v0/test_token_tmp.c2", "()[]");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	if (tokenstream_next(ts) != TOKEN_PARENT_OPEN) fail("expected TOKEN_PARENT_OPEN");
 	if (tokenstream_next(ts) != TOKEN_PARENT_CLOSE) fail("expected TOKEN_PARENT_CLOSE");
@@ -83,8 +95,11 @@ static void test_tokenstream_parentheses_and_brackets(void) {
 
 static void test_tokenstream_comma(void) {
 	write_test_file("v0/test_token_tmp.c2", "a,b,c");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	if (tokenstream_next(ts) != TOKEN_IDENTIFIER) fail("expected a");
 	if (tokenstream_next(ts) != TOKEN_COMMA) fail("expected comma");
@@ -98,8 +113,11 @@ static void test_tokenstream_comma(void) {
 
 static void test_tokenstream_whitespace_ignored(void) {
 	write_test_file("v0/test_token_tmp.c2", "  module  \n\t import  ;  ");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	if (tokenstream_next(ts) != TOKEN_MODULE) fail("expected TOKEN_MODULE");
 	if (tokenstream_next(ts) != TOKEN_IMPORT) fail("expected TOKEN_IMPORT");
@@ -111,8 +129,11 @@ static void test_tokenstream_whitespace_ignored(void) {
 
 static void test_tokenstream_void_function_signature(void) {
 	write_test_file("v0/test_token_tmp.c2", "void main()");
-	TokenStream* ts = tokenstream_open("v0/test_token_tmp.c2");
-	if (ts == NULL) fail("could not open file");
+	Buffer* buf = buffer_open_file("v0/test_token_tmp.c2");
+	if (buf == NULL) fail("could not open file");
+
+	TokenStream* ts = tokenstream_open(buf);
+	if (ts == NULL) fail("could not create tokenstream");
 
 	if (tokenstream_next(ts) != TOKEN_VOID) fail("expected TOKEN_VOID");
 	if (tokenstream_next(ts) != TOKEN_IDENTIFIER) fail("expected TOKEN_IDENTIFIER");
diff --git a/v0/token.c b/v0/token.c
index ce3a6ae..d746055 100644
--- a/v0/token.c
+++ b/v0/token.c
@@ -103,19 +103,15 @@ static Token read_keyword_or_identifier(TokenStream* ts, char first) {
 	return TOKEN_IDENTIFIER;
 }
 
-TokenStream* tokenstream_open(const char* path) {
-	if (path == NULL) return NULL;
-
-	Buffer* buf = buffer_open_file(path);
-	if (buf == NULL) return NULL;
+TokenStream* tokenstream_open(Buffer* buffer) {
+	if (buffer == NULL) return NULL;
 
 	TokenStream* ts = (TokenStream*)malloc(sizeof(struct TokenStream));
 	if (ts == NULL) {
-		buffer_close(buf);
 		return NULL;
 	}
 
-	ts->buffer = buf;
+	ts->buffer = buffer;
 	ts->lookahead = 0;
 	ts->has_lookahead = 0;
 	return ts;
diff --git a/v0/token.h b/v0/token.h
index 01f4dca..1fb17ed 100644
--- a/v0/token.h
+++ b/v0/token.h
@@ -4,6 +4,8 @@
 #ifndef TOKEN_H
 #define TOKEN_H
 
+#include "buffer.h"
+
 typedef enum {
 	// Keywords
 	TOKEN_MODULE,
@@ -34,11 +36,14 @@ typedef struct TokenStream TokenStream;
 const char* token_to_string(Token token);
 
 /**
- * Opens a file and returns a TokenStream for it.
- * @param path The path to the file.
+ * Returns a TokenStream for a given buffer.
+ * 
+ * When the tokenstream is closed, the underlying buffer is also closed.
+ * 
+ * @param buffer The buffer to read from.
  * @returns A handle to the TokenStream.
  */
-TokenStream* tokenstream_open(const char* path);
+TokenStream* tokenstream_open(Buffer* buffer);
 
 /**
  * Closes a TokenStream.