In most cases, the state boundary line controls.
For example, Kentucky owns the Ohio River along its border with Ohio, Indiana, and Illinois. In essence, the boundary between Kentucky and these three future states is the low point of the Ohio River's northernmost bank. This was established by a Congressionally passed law in 1792.
The boundary is defined by federal statute on a case by case basis for each of the rivers that is a boundary between U.S. states. Most codifications of a state's statutory laws also restate the federal statute defining the state's boundary at the beginning before the actual portion setting forth state laws begins.
In the case of the Rio Grande river, which is a boundary between the U.S. and Mexico, an international boundary treaty determines where the boundary between the U.S. and Mexico is for this purpose. International treaties similarly define all other international water boundaries of the United States.
This isn't the end of the story, however, because many activities that take place in the "navigable waters of the United States" (which basically means with respect to rivers, rivers wide and deep enough to be used for commercial transportation, like the Mississippi and Ohio Rivers and the Colorado River) are governed by federal admiralty law, rather than by state law (even if both sides of the river are in the same state).
For example, the employment law of ship crews operating on the navigable waters of the United States are governed by admiralty law, and so are most laws governing tort liability in collisions between boats/ships, and most laws governing salvage of boats/ships that have sunk.
Some admiralty law cases must be litigated in the U.S. District Courts (i.e. in the federal court system's trial courts) which have exclusive jurisdiction over some admiralty law issues, while other admiralty law issues, while still governed by federal admiralty law, may be litigated in either state or federal courts which have concurrent jurisdiction over some admiralty law issues.
Still, to the extent that admiralty law does not apply, the state law of the state in which that part of the river is located will govern.
Fun fact: As a young lawyer I once was part of a group of lawyers litigating an admiralty law case arising out of an incident taking place on the Colorado River near Grand Junction, Colorado.
Do I have any defense for breaking breaking state B's law if I assumed I was still in state A jurisdiction?
No. You could argue at trial that you really were in state A jurisdiction. But if the finder of fact at trial determines that you were in state B, then state B law applies.